diff --git a/_modules/cyto_dl/api/cyto_dl_model/cyto_dl_base_model.html b/_modules/cyto_dl/api/cyto_dl_model/cyto_dl_base_model.html
index 0a807bd73..eed56ea6e 100644
--- a/_modules/cyto_dl/api/cyto_dl_model/cyto_dl_base_model.html
+++ b/_modules/cyto_dl/api/cyto_dl_model/cyto_dl_base_model.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/api/cyto_dl_model/segmentation_plugin_model.html b/_modules/cyto_dl/api/cyto_dl_model/segmentation_plugin_model.html
index 3b20387b4..0c3250b89 100644
--- a/_modules/cyto_dl/api/cyto_dl_model/segmentation_plugin_model.html
+++ b/_modules/cyto_dl/api/cyto_dl_model/segmentation_plugin_model.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/api/data.html b/_modules/cyto_dl/api/data.html
index d6f227c4f..ef0d30422 100644
--- a/_modules/cyto_dl/api/data.html
+++ b/_modules/cyto_dl/api/data.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/api/model.html b/_modules/cyto_dl/api/model.html
index cead51d1b..524c2571e 100644
--- a/_modules/cyto_dl/api/model.html
+++ b/_modules/cyto_dl/api/model.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/callbacks/csv_saver.html b/_modules/cyto_dl/callbacks/csv_saver.html
index 7e4ae6b53..7989b1566 100644
--- a/_modules/cyto_dl/callbacks/csv_saver.html
+++ b/_modules/cyto_dl/callbacks/csv_saver.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/callbacks/image_saver.html b/_modules/cyto_dl/callbacks/image_saver.html
index dc86321d2..530dfa8d2 100644
--- a/_modules/cyto_dl/callbacks/image_saver.html
+++ b/_modules/cyto_dl/callbacks/image_saver.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/callbacks/layer_freeze.html b/_modules/cyto_dl/callbacks/layer_freeze.html
index 7596e5d88..6d76992a8 100644
--- a/_modules/cyto_dl/callbacks/layer_freeze.html
+++ b/_modules/cyto_dl/callbacks/layer_freeze.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/callbacks/model_utils.html b/_modules/cyto_dl/callbacks/model_utils.html
index c72e2b853..3b0e514f5 100644
--- a/_modules/cyto_dl/callbacks/model_utils.html
+++ b/_modules/cyto_dl/callbacks/model_utils.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/callbacks/outlier_detection.html b/_modules/cyto_dl/callbacks/outlier_detection.html
index 404902c5f..4adf74820 100644
--- a/_modules/cyto_dl/callbacks/outlier_detection.html
+++ b/_modules/cyto_dl/callbacks/outlier_detection.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/dataframe/readers.html b/_modules/cyto_dl/dataframe/readers.html
index 4616bfb63..3fb18895a 100644
--- a/_modules/cyto_dl/dataframe/readers.html
+++ b/_modules/cyto_dl/dataframe/readers.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/dataframe/transforms/filter.html b/_modules/cyto_dl/dataframe/transforms/filter.html
index 0b111a2b3..5ec47e560 100644
--- a/_modules/cyto_dl/dataframe/transforms/filter.html
+++ b/_modules/cyto_dl/dataframe/transforms/filter.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/dataframe/transforms/group_cols.html b/_modules/cyto_dl/dataframe/transforms/group_cols.html
index 56370dd25..7386f76a3 100644
--- a/_modules/cyto_dl/dataframe/transforms/group_cols.html
+++ b/_modules/cyto_dl/dataframe/transforms/group_cols.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/dataframe/transforms/misc.html b/_modules/cyto_dl/dataframe/transforms/misc.html
index 86c0b125e..fc1a00d02 100644
--- a/_modules/cyto_dl/dataframe/transforms/misc.html
+++ b/_modules/cyto_dl/dataframe/transforms/misc.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/dataframe/transforms/split.html b/_modules/cyto_dl/dataframe/transforms/split.html
index 93de69617..3f9438919 100644
--- a/_modules/cyto_dl/dataframe/transforms/split.html
+++ b/_modules/cyto_dl/dataframe/transforms/split.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/datamodules/array.html b/_modules/cyto_dl/datamodules/array.html
index b9841f874..29ba62586 100644
--- a/_modules/cyto_dl/datamodules/array.html
+++ b/_modules/cyto_dl/datamodules/array.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/datamodules/data_dict.html b/_modules/cyto_dl/datamodules/data_dict.html
index 7971dadda..e87b46a15 100644
--- a/_modules/cyto_dl/datamodules/data_dict.html
+++ b/_modules/cyto_dl/datamodules/data_dict.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/datamodules/dataframe/dataframe_datamodule.html b/_modules/cyto_dl/datamodules/dataframe/dataframe_datamodule.html
index 47e60458f..9e0f8f00e 100644
--- a/_modules/cyto_dl/datamodules/dataframe/dataframe_datamodule.html
+++ b/_modules/cyto_dl/datamodules/dataframe/dataframe_datamodule.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/datamodules/dataframe/grouped_dataframe_datamodule.html b/_modules/cyto_dl/datamodules/dataframe/grouped_dataframe_datamodule.html
index 0826cc03a..9451bb177 100644
--- a/_modules/cyto_dl/datamodules/dataframe/grouped_dataframe_datamodule.html
+++ b/_modules/cyto_dl/datamodules/dataframe/grouped_dataframe_datamodule.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/datamodules/dataframe/utils.html b/_modules/cyto_dl/datamodules/dataframe/utils.html
index a7a089fd6..e91de9c7f 100644
--- a/_modules/cyto_dl/datamodules/dataframe/utils.html
+++ b/_modules/cyto_dl/datamodules/dataframe/utils.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/datamodules/folder.html b/_modules/cyto_dl/datamodules/folder.html
index bbe1ddcbd..f804443c9 100644
--- a/_modules/cyto_dl/datamodules/folder.html
+++ b/_modules/cyto_dl/datamodules/folder.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/datamodules/multidim_image.html b/_modules/cyto_dl/datamodules/multidim_image.html
index 7da3eea18..c4cdb345a 100644
--- a/_modules/cyto_dl/datamodules/multidim_image.html
+++ b/_modules/cyto_dl/datamodules/multidim_image.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/datamodules/smartcache.html b/_modules/cyto_dl/datamodules/smartcache.html
index 738c29ea1..0e2b8112b 100644
--- a/_modules/cyto_dl/datamodules/smartcache.html
+++ b/_modules/cyto_dl/datamodules/smartcache.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/eval.html b/_modules/cyto_dl/eval.html
index 5cf781ab6..56e509d39 100644
--- a/_modules/cyto_dl/eval.html
+++ b/_modules/cyto_dl/eval.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/image/io/aicsimage_loader.html b/_modules/cyto_dl/image/io/aicsimage_loader.html
index c1244c360..08cedc264 100644
--- a/_modules/cyto_dl/image/io/aicsimage_loader.html
+++ b/_modules/cyto_dl/image/io/aicsimage_loader.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/image/io/monai_bio_reader.html b/_modules/cyto_dl/image/io/monai_bio_reader.html
index ecc195e18..964009579 100644
--- a/_modules/cyto_dl/image/io/monai_bio_reader.html
+++ b/_modules/cyto_dl/image/io/monai_bio_reader.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/image/io/numpy_reader.html b/_modules/cyto_dl/image/io/numpy_reader.html
index 86f15adbf..36500463e 100644
--- a/_modules/cyto_dl/image/io/numpy_reader.html
+++ b/_modules/cyto_dl/image/io/numpy_reader.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/image/io/ome_zarr_reader.html b/_modules/cyto_dl/image/io/ome_zarr_reader.html
index 3c7b18769..c5aa9b5f8 100644
--- a/_modules/cyto_dl/image/io/ome_zarr_reader.html
+++ b/_modules/cyto_dl/image/io/ome_zarr_reader.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/image/io/polygon_loader.html b/_modules/cyto_dl/image/io/polygon_loader.html
index 5c64ed987..9d4587379 100644
--- a/_modules/cyto_dl/image/io/polygon_loader.html
+++ b/_modules/cyto_dl/image/io/polygon_loader.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/image/io/skimage_reader.html b/_modules/cyto_dl/image/io/skimage_reader.html
index fcadde366..b7c5ad815 100644
--- a/_modules/cyto_dl/image/io/skimage_reader.html
+++ b/_modules/cyto_dl/image/io/skimage_reader.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/image/transforms/add_meta.html b/_modules/cyto_dl/image/transforms/add_meta.html
index d2a87e142..9e38e3bfe 100644
--- a/_modules/cyto_dl/image/transforms/add_meta.html
+++ b/_modules/cyto_dl/image/transforms/add_meta.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/image/transforms/bright_sampler.html b/_modules/cyto_dl/image/transforms/bright_sampler.html
index 927e2430c..c021543f2 100644
--- a/_modules/cyto_dl/image/transforms/bright_sampler.html
+++ b/_modules/cyto_dl/image/transforms/bright_sampler.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/image/transforms/clip.html b/_modules/cyto_dl/image/transforms/clip.html
index 69e65006f..f25126deb 100644
--- a/_modules/cyto_dl/image/transforms/clip.html
+++ b/_modules/cyto_dl/image/transforms/clip.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/image/transforms/contrastadjust.html b/_modules/cyto_dl/image/transforms/contrastadjust.html
index c8e8418f2..bb84247dc 100644
--- a/_modules/cyto_dl/image/transforms/contrastadjust.html
+++ b/_modules/cyto_dl/image/transforms/contrastadjust.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/image/transforms/generate_jepa_masks.html b/_modules/cyto_dl/image/transforms/generate_jepa_masks.html
index 2a691822d..cba5a1e40 100644
--- a/_modules/cyto_dl/image/transforms/generate_jepa_masks.html
+++ b/_modules/cyto_dl/image/transforms/generate_jepa_masks.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
@@ -440,6 +447,8 @@ <h1>Source code for cyto_dl.image.transforms.generate_jepa_masks</h1><div class=
 <span class="kn">from</span> <span class="nn">monai.transforms</span> <span class="kn">import</span> <span class="n">RandomizableTransform</span>
 <span class="kn">from</span> <span class="nn">skimage.segmentation</span> <span class="kn">import</span> <span class="n">find_boundaries</span>
 
+<span class="kn">from</span> <span class="nn">cyto_dl.nn.vits.utils</span> <span class="kn">import</span> <span class="n">match_tuple_dimensions</span>
+
 
 <div class="viewcode-block" id="JEPAMaskGenerator"><a class="viewcode-back" href="../../../../cyto_dl.image.transforms.generate_jepa_masks.html#cyto_dl.image.transforms.generate_jepa_masks.JEPAMaskGenerator">[docs]</a><span class="k">class</span> <span class="nc">JEPAMaskGenerator</span><span class="p">(</span><span class="n">RandomizableTransform</span><span class="p">):</span>
 <span class="w">    </span><span class="sd">&quot;&quot;&quot;Transform for generating Block-contiguous masks for JEPA training.</span>
@@ -450,6 +459,7 @@ <h1>Source code for cyto_dl.image.transforms.generate_jepa_masks</h1><div class=
 
     <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span>
         <span class="bp">self</span><span class="p">,</span>
+        <span class="n">spatial_dims</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span>
         <span class="n">mask_size</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">12</span><span class="p">,</span>
         <span class="n">block_aspect_ratio</span><span class="p">:</span> <span class="n">Tuple</span><span class="p">[</span><span class="nb">float</span><span class="p">]</span> <span class="o">=</span> <span class="p">(</span><span class="mf">0.5</span><span class="p">,</span> <span class="mf">1.5</span><span class="p">),</span>
         <span class="n">num_patches</span><span class="p">:</span> <span class="n">Tuple</span><span class="p">[</span><span class="nb">float</span><span class="p">]</span> <span class="o">=</span> <span class="p">(</span><span class="mi">6</span><span class="p">,</span> <span class="mi">24</span><span class="p">,</span> <span class="mi">24</span><span class="p">),</span>
@@ -458,6 +468,8 @@ <h1>Source code for cyto_dl.image.transforms.generate_jepa_masks</h1><div class=
 <span class="w">        </span><span class="sd">&quot;&quot;&quot;</span>
 <span class="sd">        Parameters</span>
 <span class="sd">        ----------</span>
+<span class="sd">        spatial_dims : int</span>
+<span class="sd">            The number of spatial dimensions of the image (2 or 3)</span>
 <span class="sd">        mask_size : int, optional</span>
 <span class="sd">            The size of the blocks used to generate mask. Block dimensions are determined by the mask size and an aspect ratio sampled from the range  `block_aspect_ratio`</span>
 <span class="sd">        block_aspect_ratio : Tuple[float], optional</span>
@@ -467,7 +479,9 @@ <h1>Source code for cyto_dl.image.transforms.generate_jepa_masks</h1><div class=
 <span class="sd">        mask_ratio : float, optional</span>
 <span class="sd">            The proportion of the image to be masked</span>
 <span class="sd">        &quot;&quot;&quot;</span>
-        <span class="k">assert</span> <span class="n">mask_ratio</span> <span class="o">&lt;</span> <span class="mi">1</span><span class="p">,</span> <span class="s2">&quot;mask_ratio must be less than 1&quot;</span>
+        <span class="k">assert</span> <span class="mi">0</span> <span class="o">&lt;</span> <span class="n">mask_ratio</span> <span class="o">&lt;</span> <span class="mi">1</span><span class="p">,</span> <span class="s2">&quot;mask_ratio must be between 0 and 1&quot;</span>
+
+        <span class="n">num_patches</span> <span class="o">=</span> <span class="n">match_tuple_dimensions</span><span class="p">(</span><span class="n">spatial_dims</span><span class="p">,</span> <span class="p">[</span><span class="n">num_patches</span><span class="p">])[</span><span class="mi">0</span><span class="p">]</span>
         <span class="k">assert</span> <span class="n">mask_size</span> <span class="o">*</span> <span class="nb">max</span><span class="p">(</span><span class="n">block_aspect_ratio</span><span class="p">)</span> <span class="o">&lt;</span> <span class="nb">min</span><span class="p">(</span>
             <span class="n">num_patches</span><span class="p">[</span><span class="o">-</span><span class="mi">2</span><span class="p">:]</span>
         <span class="p">),</span> <span class="s2">&quot;mask_size * max mask aspect ratio must be less than the smallest dimension of num_patches&quot;</span>
@@ -481,7 +495,7 @@ <h1>Source code for cyto_dl.image.transforms.generate_jepa_masks</h1><div class=
         <span class="bp">self</span><span class="o">.</span><span class="n">mask</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="n">num_patches</span><span class="p">)</span>
         <span class="bp">self</span><span class="o">.</span><span class="n">edge_mask</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">ones</span><span class="p">(</span><span class="n">num_patches</span><span class="p">)</span>
 
-        <span class="bp">self</span><span class="o">.</span><span class="n">spatial_dims</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="n">num_patches</span><span class="p">)</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">spatial_dims</span> <span class="o">=</span> <span class="n">spatial_dims</span>
         <span class="c1"># create a mask that identified pixels on the edge of the image</span>
         <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">spatial_dims</span> <span class="o">==</span> <span class="mi">3</span><span class="p">:</span>
             <span class="bp">self</span><span class="o">.</span><span class="n">edge_mask</span><span class="p">[</span><span class="mi">1</span><span class="p">:</span><span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">:</span><span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">:</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span> <span class="o">=</span> <span class="mi">0</span>
diff --git a/_modules/cyto_dl/image/transforms/merge.html b/_modules/cyto_dl/image/transforms/merge.html
index e5c7888e8..72fa990bf 100644
--- a/_modules/cyto_dl/image/transforms/merge.html
+++ b/_modules/cyto_dl/image/transforms/merge.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/image/transforms/multiscale_cropper.html b/_modules/cyto_dl/image/transforms/multiscale_cropper.html
index 847814bd4..55ecf7775 100644
--- a/_modules/cyto_dl/image/transforms/multiscale_cropper.html
+++ b/_modules/cyto_dl/image/transforms/multiscale_cropper.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/image/transforms/pad.html b/_modules/cyto_dl/image/transforms/pad.html
index 421422979..c1fa35300 100644
--- a/_modules/cyto_dl/image/transforms/pad.html
+++ b/_modules/cyto_dl/image/transforms/pad.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/image/transforms/project.html b/_modules/cyto_dl/image/transforms/project.html
index 6b93e89bc..72255f61d 100644
--- a/_modules/cyto_dl/image/transforms/project.html
+++ b/_modules/cyto_dl/image/transforms/project.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/image/transforms/save.html b/_modules/cyto_dl/image/transforms/save.html
index 8a4de7fb2..33a00afc5 100644
--- a/_modules/cyto_dl/image/transforms/save.html
+++ b/_modules/cyto_dl/image/transforms/save.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/image/transforms/track_transforms.html b/_modules/cyto_dl/image/transforms/track_transforms.html
index 91bdfd6aa..8a2cb8a0f 100644
--- a/_modules/cyto_dl/image/transforms/track_transforms.html
+++ b/_modules/cyto_dl/image/transforms/track_transforms.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/loggers/mlflow.html b/_modules/cyto_dl/loggers/mlflow.html
index df68a6f51..3525b73a7 100644
--- a/_modules/cyto_dl/loggers/mlflow.html
+++ b/_modules/cyto_dl/loggers/mlflow.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/models/base_model.html b/_modules/cyto_dl/models/base_model.html
index 5941698e7..3812db233 100644
--- a/_modules/cyto_dl/models/base_model.html
+++ b/_modules/cyto_dl/models/base_model.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/models/basic_model.html b/_modules/cyto_dl/models/basic_model.html
index e5497b224..f963af6d4 100644
--- a/_modules/cyto_dl/models/basic_model.html
+++ b/_modules/cyto_dl/models/basic_model.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/models/classification/classification.html b/_modules/cyto_dl/models/classification/classification.html
index c21402f44..bdb38664f 100644
--- a/_modules/cyto_dl/models/classification/classification.html
+++ b/_modules/cyto_dl/models/classification/classification.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/models/classification/timepoint_classification.html b/_modules/cyto_dl/models/classification/timepoint_classification.html
index 6dd55e9b7..c59801e9c 100644
--- a/_modules/cyto_dl/models/classification/timepoint_classification.html
+++ b/_modules/cyto_dl/models/classification/timepoint_classification.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/models/contrastive/contrastive.html b/_modules/cyto_dl/models/contrastive/contrastive.html
index 4117e051b..6c869917a 100644
--- a/_modules/cyto_dl/models/contrastive/contrastive.html
+++ b/_modules/cyto_dl/models/contrastive/contrastive.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/models/im2im/gan.html b/_modules/cyto_dl/models/im2im/gan.html
index 7cf89010b..f6c84d3a7 100644
--- a/_modules/cyto_dl/models/im2im/gan.html
+++ b/_modules/cyto_dl/models/im2im/gan.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/models/im2im/multi_task.html b/_modules/cyto_dl/models/im2im/multi_task.html
index 0ea5330fa..2b635fa34 100644
--- a/_modules/cyto_dl/models/im2im/multi_task.html
+++ b/_modules/cyto_dl/models/im2im/multi_task.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/models/im2im/utils/instance_seg.html b/_modules/cyto_dl/models/im2im/utils/instance_seg.html
index b443b8824..9b9209ca5 100644
--- a/_modules/cyto_dl/models/im2im/utils/instance_seg.html
+++ b/_modules/cyto_dl/models/im2im/utils/instance_seg.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/models/im2im/utils/noise_annealer.html b/_modules/cyto_dl/models/im2im/utils/noise_annealer.html
index 6df63a4df..dae8b3678 100644
--- a/_modules/cyto_dl/models/im2im/utils/noise_annealer.html
+++ b/_modules/cyto_dl/models/im2im/utils/noise_annealer.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/models/im2im/utils/postprocessing/act_thresh_label.html b/_modules/cyto_dl/models/im2im/utils/postprocessing/act_thresh_label.html
index 1d919d05d..7bcfdf88f 100644
--- a/_modules/cyto_dl/models/im2im/utils/postprocessing/act_thresh_label.html
+++ b/_modules/cyto_dl/models/im2im/utils/postprocessing/act_thresh_label.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/models/im2im/utils/postprocessing/auto_thresh.html b/_modules/cyto_dl/models/im2im/utils/postprocessing/auto_thresh.html
index e64c15b4c..418d30d9a 100644
--- a/_modules/cyto_dl/models/im2im/utils/postprocessing/auto_thresh.html
+++ b/_modules/cyto_dl/models/im2im/utils/postprocessing/auto_thresh.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/models/im2im/utils/postprocessing/dict_to_im.html b/_modules/cyto_dl/models/im2im/utils/postprocessing/dict_to_im.html
index 03eb2f698..14639be7b 100644
--- a/_modules/cyto_dl/models/im2im/utils/postprocessing/dict_to_im.html
+++ b/_modules/cyto_dl/models/im2im/utils/postprocessing/dict_to_im.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/models/jepa/ijepa.html b/_modules/cyto_dl/models/jepa/ijepa.html
index 4f2953547..7e038fb66 100644
--- a/_modules/cyto_dl/models/jepa/ijepa.html
+++ b/_modules/cyto_dl/models/jepa/ijepa.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/models/jepa/iwm.html b/_modules/cyto_dl/models/jepa/iwm.html
index 702a8e567..fdc046f03 100644
--- a/_modules/cyto_dl/models/jepa/iwm.html
+++ b/_modules/cyto_dl/models/jepa/iwm.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/models/jepa/jepa_base.html b/_modules/cyto_dl/models/jepa/jepa_base.html
index c83810931..6d653d26c 100644
--- a/_modules/cyto_dl/models/jepa/jepa_base.html
+++ b/_modules/cyto_dl/models/jepa/jepa_base.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/models/utils/mlflow.html b/_modules/cyto_dl/models/utils/mlflow.html
index 543a3683e..d51d83214 100644
--- a/_modules/cyto_dl/models/utils/mlflow.html
+++ b/_modules/cyto_dl/models/utils/mlflow.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/models/utils/utils.html b/_modules/cyto_dl/models/utils/utils.html
index fad8df404..a827d1221 100644
--- a/_modules/cyto_dl/models/utils/utils.html
+++ b/_modules/cyto_dl/models/utils/utils.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/models/vae/base_vae.html b/_modules/cyto_dl/models/vae/base_vae.html
index 051a40e39..385b4cff2 100644
--- a/_modules/cyto_dl/models/vae/base_vae.html
+++ b/_modules/cyto_dl/models/vae/base_vae.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/models/vae/priors/abstract_prior.html b/_modules/cyto_dl/models/vae/priors/abstract_prior.html
index 2ebabb901..c906c8506 100644
--- a/_modules/cyto_dl/models/vae/priors/abstract_prior.html
+++ b/_modules/cyto_dl/models/vae/priors/abstract_prior.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/models/vae/priors/gaussian.html b/_modules/cyto_dl/models/vae/priors/gaussian.html
index 91423a15c..675ab33c7 100644
--- a/_modules/cyto_dl/models/vae/priors/gaussian.html
+++ b/_modules/cyto_dl/models/vae/priors/gaussian.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/models/vae/priors/identity_prior.html b/_modules/cyto_dl/models/vae/priors/identity_prior.html
index f2703a326..dee86df55 100644
--- a/_modules/cyto_dl/models/vae/priors/identity_prior.html
+++ b/_modules/cyto_dl/models/vae/priors/identity_prior.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/models/vae/priors/joint_prior.html b/_modules/cyto_dl/models/vae/priors/joint_prior.html
index d835277b0..1c27b5d46 100644
--- a/_modules/cyto_dl/models/vae/priors/joint_prior.html
+++ b/_modules/cyto_dl/models/vae/priors/joint_prior.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/nn/discriminators/multi_scale_discriminator.html b/_modules/cyto_dl/nn/discriminators/multi_scale_discriminator.html
index 80fa60d08..f76cdd39a 100644
--- a/_modules/cyto_dl/nn/discriminators/multi_scale_discriminator.html
+++ b/_modules/cyto_dl/nn/discriminators/multi_scale_discriminator.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/nn/discriminators/n_layer_discriminator.html b/_modules/cyto_dl/nn/discriminators/n_layer_discriminator.html
index a43867db5..116c0cc38 100644
--- a/_modules/cyto_dl/nn/discriminators/n_layer_discriminator.html
+++ b/_modules/cyto_dl/nn/discriminators/n_layer_discriminator.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/nn/head/base_head.html b/_modules/cyto_dl/nn/head/base_head.html
index 07a7a74f4..f92c1386d 100644
--- a/_modules/cyto_dl/nn/head/base_head.html
+++ b/_modules/cyto_dl/nn/head/base_head.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/nn/head/gan_head.html b/_modules/cyto_dl/nn/head/gan_head.html
index 1b3193278..ad9a1f193 100644
--- a/_modules/cyto_dl/nn/head/gan_head.html
+++ b/_modules/cyto_dl/nn/head/gan_head.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/nn/head/gan_head_superres.html b/_modules/cyto_dl/nn/head/gan_head_superres.html
index 1561e01ae..a0e714531 100644
--- a/_modules/cyto_dl/nn/head/gan_head_superres.html
+++ b/_modules/cyto_dl/nn/head/gan_head_superres.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/nn/head/mae_head.html b/_modules/cyto_dl/nn/head/mae_head.html
index 1e550701f..2b2c3e576 100644
--- a/_modules/cyto_dl/nn/head/mae_head.html
+++ b/_modules/cyto_dl/nn/head/mae_head.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/nn/head/mask_head.html b/_modules/cyto_dl/nn/head/mask_head.html
index b338d0c94..1e6d2fdd2 100644
--- a/_modules/cyto_dl/nn/head/mask_head.html
+++ b/_modules/cyto_dl/nn/head/mask_head.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/nn/head/res_blocks_head.html b/_modules/cyto_dl/nn/head/res_blocks_head.html
index 8e3731ab5..3ad59aaea 100644
--- a/_modules/cyto_dl/nn/head/res_blocks_head.html
+++ b/_modules/cyto_dl/nn/head/res_blocks_head.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/nn/head/vic_reg.html b/_modules/cyto_dl/nn/head/vic_reg.html
index 4a5365c62..79f7365da 100644
--- a/_modules/cyto_dl/nn/head/vic_reg.html
+++ b/_modules/cyto_dl/nn/head/vic_reg.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/nn/hr_skip.html b/_modules/cyto_dl/nn/hr_skip.html
index 240ea3d31..3a1024b8d 100644
--- a/_modules/cyto_dl/nn/hr_skip.html
+++ b/_modules/cyto_dl/nn/hr_skip.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/nn/losses/adversarial_loss.html b/_modules/cyto_dl/nn/losses/adversarial_loss.html
index 952595936..fc2dad903 100644
--- a/_modules/cyto_dl/nn/losses/adversarial_loss.html
+++ b/_modules/cyto_dl/nn/losses/adversarial_loss.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/nn/losses/chamfer_loss.html b/_modules/cyto_dl/nn/losses/chamfer_loss.html
index 4005d0f49..366032fa1 100644
--- a/_modules/cyto_dl/nn/losses/chamfer_loss.html
+++ b/_modules/cyto_dl/nn/losses/chamfer_loss.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/nn/losses/continuous_bernoulli.html b/_modules/cyto_dl/nn/losses/continuous_bernoulli.html
index 408d38cf5..a754daed0 100644
--- a/_modules/cyto_dl/nn/losses/continuous_bernoulli.html
+++ b/_modules/cyto_dl/nn/losses/continuous_bernoulli.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/nn/losses/cosine_loss.html b/_modules/cyto_dl/nn/losses/cosine_loss.html
index db3bb5910..2cdfb69c9 100644
--- a/_modules/cyto_dl/nn/losses/cosine_loss.html
+++ b/_modules/cyto_dl/nn/losses/cosine_loss.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/nn/losses/gan_loss.html b/_modules/cyto_dl/nn/losses/gan_loss.html
index 97be27afc..737b253bd 100644
--- a/_modules/cyto_dl/nn/losses/gan_loss.html
+++ b/_modules/cyto_dl/nn/losses/gan_loss.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/nn/losses/gaussian_nll_loss.html b/_modules/cyto_dl/nn/losses/gaussian_nll_loss.html
index d5e179ff2..ddaf6607f 100644
--- a/_modules/cyto_dl/nn/losses/gaussian_nll_loss.html
+++ b/_modules/cyto_dl/nn/losses/gaussian_nll_loss.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/nn/losses/loss_wrapper.html b/_modules/cyto_dl/nn/losses/loss_wrapper.html
index 44e49ca99..604cb4a2a 100644
--- a/_modules/cyto_dl/nn/losses/loss_wrapper.html
+++ b/_modules/cyto_dl/nn/losses/loss_wrapper.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/nn/losses/threshold_loss.html b/_modules/cyto_dl/nn/losses/threshold_loss.html
index 3c4f17e87..3ff9a904b 100644
--- a/_modules/cyto_dl/nn/losses/threshold_loss.html
+++ b/_modules/cyto_dl/nn/losses/threshold_loss.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/nn/losses/vic_reg.html b/_modules/cyto_dl/nn/losses/vic_reg.html
index 1f9572fb1..7d9cd4dc2 100644
--- a/_modules/cyto_dl/nn/losses/vic_reg.html
+++ b/_modules/cyto_dl/nn/losses/vic_reg.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/nn/losses/weibull.html b/_modules/cyto_dl/nn/losses/weibull.html
index db7a8477c..4e3121e6b 100644
--- a/_modules/cyto_dl/nn/losses/weibull.html
+++ b/_modules/cyto_dl/nn/losses/weibull.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/nn/losses/weighted_mse_loss.html b/_modules/cyto_dl/nn/losses/weighted_mse_loss.html
index 11c277f19..0201880ce 100644
--- a/_modules/cyto_dl/nn/losses/weighted_mse_loss.html
+++ b/_modules/cyto_dl/nn/losses/weighted_mse_loss.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/nn/mlp.html b/_modules/cyto_dl/nn/mlp.html
index df3632774..9f95779e0 100644
--- a/_modules/cyto_dl/nn/mlp.html
+++ b/_modules/cyto_dl/nn/mlp.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/nn/res_unit.html b/_modules/cyto_dl/nn/res_unit.html
index 2f1ed8a2f..0331dd33d 100644
--- a/_modules/cyto_dl/nn/res_unit.html
+++ b/_modules/cyto_dl/nn/res_unit.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/nn/spatial_transformer.html b/_modules/cyto_dl/nn/spatial_transformer.html
index e41782e5b..067235aea 100644
--- a/_modules/cyto_dl/nn/spatial_transformer.html
+++ b/_modules/cyto_dl/nn/spatial_transformer.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/nn/track_sequence_predictor.html b/_modules/cyto_dl/nn/track_sequence_predictor.html
index 7401c3a57..299afa60d 100644
--- a/_modules/cyto_dl/nn/track_sequence_predictor.html
+++ b/_modules/cyto_dl/nn/track_sequence_predictor.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/nn/vits/blocks/cross_attention.html b/_modules/cyto_dl/nn/vits/blocks/cross_attention.html
index b5cfde889..eebc662cf 100644
--- a/_modules/cyto_dl/nn/vits/blocks/cross_attention.html
+++ b/_modules/cyto_dl/nn/vits/blocks/cross_attention.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/nn/vits/blocks/intermediate_weigher.html b/_modules/cyto_dl/nn/vits/blocks/intermediate_weigher.html
index c2de9d275..c5a3a9ebc 100644
--- a/_modules/cyto_dl/nn/vits/blocks/intermediate_weigher.html
+++ b/_modules/cyto_dl/nn/vits/blocks/intermediate_weigher.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/nn/vits/blocks/patchify.html b/_modules/cyto_dl/nn/vits/blocks/masked_unit_attention.html
similarity index 60%
rename from _modules/cyto_dl/nn/vits/blocks/patchify.html
rename to _modules/cyto_dl/nn/vits/blocks/masked_unit_attention.html
index e6ebef2b7..0e0e192f9 100644
--- a/_modules/cyto_dl/nn/vits/blocks/patchify.html
+++ b/_modules/cyto_dl/nn/vits/blocks/masked_unit_attention.html
@@ -5,7 +5,7 @@
     <meta name="color-scheme" content="light dark"><link rel="index" title="Index" href="../../../../../genindex.html" /><link rel="search" title="Search" href="../../../../../search.html" />
 
     <!-- Generated with Sphinx 6.1.3 and Furo 2023.03.23 -->
-        <title>cyto_dl.nn.vits.blocks.patchify - cyto-dl 0.4.0 documentation</title>
+        <title>cyto_dl.nn.vits.blocks.masked_unit_attention - cyto-dl 0.4.0 documentation</title>
       <link rel="stylesheet" type="text/css" href="../../../../../_static/pygments.css" />
     <link rel="stylesheet" type="text/css" href="../../../../../_static/styles/furo.css?digest=fad236701ea90a88636c2a8c73b44ae642ed2a53" />
     <link rel="stylesheet" type="text/css" href="../../../../../_static/styles/furo-extensions.css?digest=30d1aed668e5c3a91c3e3bf6a60b675221979f0e" />
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
@@ -432,168 +439,236 @@
           </label>
         </div>
         <article role="main">
-          <h1>Source code for cyto_dl.nn.vits.blocks.patchify</h1><div class="highlight"><pre>
-<span></span><span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">List</span><span class="p">,</span> <span class="n">Optional</span>
+          <h1>Source code for cyto_dl.nn.vits.blocks.masked_unit_attention</h1><div class="highlight"><pre>
+<span></span><span class="c1"># inspired by https://github.com/facebookresearch/hiera/tree/main</span>
+<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">List</span>
 
 <span class="kn">import</span> <span class="nn">numpy</span> <span class="k">as</span> <span class="nn">np</span>
 <span class="kn">import</span> <span class="nn">torch</span>
 <span class="kn">import</span> <span class="nn">torch.nn</span> <span class="k">as</span> <span class="nn">nn</span>
-<span class="kn">from</span> <span class="nn">einops.layers.torch</span> <span class="kn">import</span> <span class="n">Rearrange</span><span class="p">,</span> <span class="n">Reduce</span>
-<span class="kn">from</span> <span class="nn">timm.models.layers</span> <span class="kn">import</span> <span class="n">trunc_normal_</span>
-
-<span class="kn">from</span> <span class="nn">cyto_dl.nn.vits.utils</span> <span class="kn">import</span> <span class="n">get_positional_embedding</span><span class="p">,</span> <span class="n">take_indexes</span>
-
-
-<div class="viewcode-block" id="random_indexes"><a class="viewcode-back" href="../../../../../cyto_dl.nn.vits.blocks.patchify.html#cyto_dl.nn.vits.blocks.patchify.random_indexes">[docs]</a><span class="k">def</span> <span class="nf">random_indexes</span><span class="p">(</span><span class="n">size</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">device</span><span class="p">):</span>
-    <span class="n">forward_indexes</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">randperm</span><span class="p">(</span><span class="n">size</span><span class="p">,</span> <span class="n">device</span><span class="o">=</span><span class="n">device</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">torch</span><span class="o">.</span><span class="n">long</span><span class="p">)</span>
-    <span class="n">backward_indexes</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">argsort</span><span class="p">(</span><span class="n">forward_indexes</span><span class="p">)</span>
-    <span class="k">return</span> <span class="n">forward_indexes</span><span class="p">,</span> <span class="n">backward_indexes</span></div>
+<span class="kn">import</span> <span class="nn">torch.nn.functional</span> <span class="k">as</span> <span class="nn">F</span>
+<span class="kn">from</span> <span class="nn">einops</span> <span class="kn">import</span> <span class="n">rearrange</span><span class="p">,</span> <span class="n">reduce</span>
+<span class="kn">from</span> <span class="nn">einops.layers.torch</span> <span class="kn">import</span> <span class="n">Reduce</span>
+<span class="kn">from</span> <span class="nn">timm.models.layers</span> <span class="kn">import</span> <span class="n">DropPath</span><span class="p">,</span> <span class="n">Mlp</span>
 
+<span class="kn">from</span> <span class="nn">cyto_dl.nn.vits.utils</span> <span class="kn">import</span> <span class="n">match_tuple_dimensions</span>
 
-<div class="viewcode-block" id="Patchify"><a class="viewcode-back" href="../../../../../cyto_dl.nn.vits.blocks.patchify.html#cyto_dl.nn.vits.blocks.patchify.Patchify">[docs]</a><span class="k">class</span> <span class="nc">Patchify</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
-<span class="w">    </span><span class="sd">&quot;&quot;&quot;Class for converting images to a masked sequence of patches with positional embeddings.&quot;&quot;&quot;</span>
 
+<div class="viewcode-block" id="MaskUnitAttention"><a class="viewcode-back" href="../../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html#cyto_dl.nn.vits.blocks.masked_unit_attention.MaskUnitAttention">[docs]</a><span class="k">class</span> <span class="nc">MaskUnitAttention</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
     <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span>
         <span class="bp">self</span><span class="p">,</span>
-        <span class="n">patch_size</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span>
-        <span class="n">emb_dim</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span>
-        <span class="n">n_patches</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span>
+        <span class="n">dim</span><span class="p">,</span>
+        <span class="n">dim_out</span><span class="p">,</span>
         <span class="n">spatial_dims</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">3</span><span class="p">,</span>
-        <span class="n">context_pixels</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">],</span>
-        <span class="n">input_channels</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">1</span><span class="p">,</span>
-        <span class="n">tasks</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">List</span><span class="p">[</span><span class="nb">str</span><span class="p">]]</span> <span class="o">=</span> <span class="p">[],</span>
-        <span class="n">learnable_pos_embedding</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">True</span><span class="p">,</span>
+        <span class="n">num_heads</span><span class="o">=</span><span class="mi">8</span><span class="p">,</span>
+        <span class="n">qkv_bias</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
+        <span class="n">attn_drop</span><span class="o">=</span><span class="mf">0.0</span><span class="p">,</span>
+        <span class="n">proj_drop</span><span class="o">=</span><span class="mf">0.0</span><span class="p">,</span>
+        <span class="n">q_stride</span><span class="o">=</span><span class="p">[</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">],</span>
+        <span class="n">patches_per_mask_unit</span><span class="o">=</span><span class="p">[</span><span class="mi">2</span><span class="p">,</span> <span class="mi">12</span><span class="p">,</span> <span class="mi">12</span><span class="p">],</span>
     <span class="p">):</span>
 <span class="w">        </span><span class="sd">&quot;&quot;&quot;</span>
 <span class="sd">        Parameters</span>
 <span class="sd">        ----------</span>
-<span class="sd">        patch_size: List[int]</span>
-<span class="sd">            Size of each patch in pix (ZYX order for 3D, YX order for 2D)</span>
-<span class="sd">        emb_dim: int</span>
-<span class="sd">            Dimension of encoder</span>
-<span class="sd">        n_patches: List[int]</span>
-<span class="sd">            Number of patches in each spatial dimension (ZYX order for 3D, YX order for 2D)</span>
-<span class="sd">        spatial_dims: int</span>
-<span class="sd">            Number of spatial dimensions</span>
-<span class="sd">        context_pixels: List[int]</span>
-<span class="sd">            Number of extra pixels around each patch to include in convolutional embedding to encoder dimension.</span>
-<span class="sd">        input_channels: int</span>
-<span class="sd">            Number of input channels</span>
-<span class="sd">        tasks: List[str]</span>
-<span class="sd">            List of tasks to encode</span>
-<span class="sd">        learnable_pos_embedding: bool</span>
-<span class="sd">            If True, learnable positional embeddings are used. If False, fixed sin/cos positional embeddings. Empirically, fixed positional embeddings work better for brightfield images.</span>
+<span class="sd">        dim : int</span>
+<span class="sd">            Dimension of the input features.</span>
+<span class="sd">        dim_out : int</span>
+<span class="sd">            Dimension of the output features.</span>
+<span class="sd">        spatial_dims : int, optional</span>
+<span class="sd">            Number of spatial dimensions, by default 3.</span>
+<span class="sd">        num_heads : int, optional</span>
+<span class="sd">            Number of attention heads, by default 8.</span>
+<span class="sd">        qkv_bias : bool, optional</span>
+<span class="sd">            If True, add a learnable bias to query, key, value, by default False.</span>
+<span class="sd">        attn_drop : float, optional</span>
+<span class="sd">            Dropout rate for attention, by default 0.0.</span>
+<span class="sd">        proj_drop : float, optional</span>
+<span class="sd">            Dropout rate for projection, by default 0.0.</span>
+<span class="sd">        q_stride : List[int], optional</span>
+<span class="sd">            Stride for query, by default [1, 1, 1].</span>
+<span class="sd">        patches_per_mask_unit : List[int], optional</span>
+<span class="sd">            Number of patches per mask unit, by default [2, 12, 12].</span>
 <span class="sd">        &quot;&quot;&quot;</span>
         <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>
-        <span class="bp">self</span><span class="o">.</span><span class="n">n_patches</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">asarray</span><span class="p">(</span><span class="n">n_patches</span><span class="p">)</span>
+        <span class="n">q_stride</span><span class="p">,</span> <span class="n">patches_per_mask_unit</span> <span class="o">=</span> <span class="n">match_tuple_dimensions</span><span class="p">(</span>
+            <span class="n">spatial_dims</span><span class="p">,</span> <span class="p">[</span><span class="n">q_stride</span><span class="p">,</span> <span class="n">patches_per_mask_unit</span><span class="p">]</span>
+        <span class="p">)</span>
+
         <span class="bp">self</span><span class="o">.</span><span class="n">spatial_dims</span> <span class="o">=</span> <span class="n">spatial_dims</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">num_heads</span> <span class="o">=</span> <span class="n">num_heads</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">head_dim</span> <span class="o">=</span> <span class="n">dim_out</span> <span class="o">//</span> <span class="n">num_heads</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">qkv</span> <span class="o">=</span> <span class="n">nn</span><span class="o">.</span><span class="n">Linear</span><span class="p">(</span><span class="n">dim</span><span class="p">,</span> <span class="n">dim_out</span> <span class="o">*</span> <span class="mi">3</span><span class="p">,</span> <span class="n">bias</span><span class="o">=</span><span class="n">qkv_bias</span><span class="p">)</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">attn_drop</span> <span class="o">=</span> <span class="n">attn_drop</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">proj</span> <span class="o">=</span> <span class="n">nn</span><span class="o">.</span><span class="n">Linear</span><span class="p">(</span><span class="n">dim_out</span><span class="p">,</span> <span class="n">dim_out</span><span class="p">)</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">proj_drop</span> <span class="o">=</span> <span class="n">nn</span><span class="o">.</span><span class="n">Dropout</span><span class="p">(</span><span class="n">proj_drop</span><span class="p">)</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">dim_out</span> <span class="o">=</span> <span class="n">dim_out</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">q_stride</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">(</span><span class="n">q_stride</span><span class="p">)</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">pooled_patches_per_mask_unit</span> <span class="o">=</span> <span class="p">(</span>
+            <span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">(</span><span class="n">patches_per_mask_unit</span><span class="p">)</span> <span class="o">/</span> <span class="bp">self</span><span class="o">.</span><span class="n">q_stride</span>
+        <span class="p">)</span><span class="o">.</span><span class="n">astype</span><span class="p">(</span><span class="nb">int</span><span class="p">)</span>
 
-        <span class="bp">self</span><span class="o">.</span><span class="n">pos_embedding</span> <span class="o">=</span> <span class="n">get_positional_embedding</span><span class="p">(</span>
-            <span class="n">n_patches</span><span class="p">,</span> <span class="n">emb_dim</span><span class="p">,</span> <span class="n">learnable</span><span class="o">=</span><span class="n">learnable_pos_embedding</span><span class="p">,</span> <span class="n">use_cls_token</span><span class="o">=</span><span class="kc">False</span>
+<div class="viewcode-block" id="MaskUnitAttention.forward"><a class="viewcode-back" href="../../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html#cyto_dl.nn.vits.blocks.masked_unit_attention.MaskUnitAttention.forward">[docs]</a>    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x</span><span class="p">):</span>
+        <span class="c1"># project and split into q,k,v embeddings</span>
+        <span class="n">qkv</span> <span class="o">=</span> <span class="n">rearrange</span><span class="p">(</span>
+            <span class="bp">self</span><span class="o">.</span><span class="n">qkv</span><span class="p">(</span><span class="n">x</span><span class="p">),</span>
+            <span class="s2">&quot;batch num_mask_units tokens_per_mask_unit (head_dim num_heads qkv) -&gt; qkv batch num_mask_units num_heads tokens_per_mask_unit head_dim&quot;</span><span class="p">,</span>
+            <span class="n">head_dim</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">head_dim</span><span class="p">,</span>
+            <span class="n">qkv</span><span class="o">=</span><span class="mi">3</span><span class="p">,</span>
+            <span class="n">num_heads</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">num_heads</span><span class="p">,</span>
         <span class="p">)</span>
 
-        <span class="n">context_pixels</span> <span class="o">=</span> <span class="n">context_pixels</span><span class="p">[:</span><span class="n">spatial_dims</span><span class="p">]</span>
-        <span class="n">weight_size</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">asarray</span><span class="p">(</span><span class="n">patch_size</span><span class="p">)</span> <span class="o">+</span> <span class="n">np</span><span class="o">.</span><span class="n">round</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">(</span><span class="n">context_pixels</span><span class="p">)</span> <span class="o">*</span> <span class="mi">2</span><span class="p">)</span><span class="o">.</span><span class="n">astype</span><span class="p">(</span><span class="nb">int</span><span class="p">)</span>
+        <span class="n">q</span><span class="p">,</span> <span class="n">k</span><span class="p">,</span> <span class="n">v</span> <span class="o">=</span> <span class="n">qkv</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">qkv</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class="n">qkv</span><span class="p">[</span><span class="mi">2</span><span class="p">]</span>
 
-        <span class="k">if</span> <span class="n">spatial_dims</span> <span class="o">==</span> <span class="mi">3</span><span class="p">:</span>
-            <span class="bp">self</span><span class="o">.</span><span class="n">conv</span> <span class="o">=</span> <span class="n">nn</span><span class="o">.</span><span class="n">Conv3d</span><span class="p">(</span>
-                <span class="n">in_channels</span><span class="o">=</span><span class="n">input_channels</span><span class="p">,</span>
-                <span class="n">out_channels</span><span class="o">=</span><span class="n">emb_dim</span><span class="p">,</span>
-                <span class="n">kernel_size</span><span class="o">=</span><span class="n">weight_size</span><span class="p">,</span>
-                <span class="n">stride</span><span class="o">=</span><span class="n">patch_size</span><span class="p">,</span>
-                <span class="n">padding</span><span class="o">=</span><span class="n">context_pixels</span><span class="p">,</span>
-            <span class="p">)</span>
-            <span class="bp">self</span><span class="o">.</span><span class="n">img2token</span> <span class="o">=</span> <span class="n">Rearrange</span><span class="p">(</span><span class="s2">&quot;b c z y x -&gt; (z y x) b c&quot;</span><span class="p">)</span>
-            <span class="bp">self</span><span class="o">.</span><span class="n">patch2img</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Sequential</span><span class="p">(</span>
-                <span class="o">*</span><span class="p">[</span>
-                    <span class="n">Rearrange</span><span class="p">(</span>
-                        <span class="s2">&quot;(n_patch_z n_patch_y n_patch_x) b c -&gt;  b c n_patch_z n_patch_y n_patch_x&quot;</span><span class="p">,</span>
-                        <span class="n">n_patch_z</span><span class="o">=</span><span class="n">n_patches</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span>
-                        <span class="n">n_patch_y</span><span class="o">=</span><span class="n">n_patches</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span>
-                        <span class="n">n_patch_x</span><span class="o">=</span><span class="n">n_patches</span><span class="p">[</span><span class="mi">2</span><span class="p">],</span>
-                    <span class="p">),</span>
-                    <span class="n">Reduce</span><span class="p">(</span>
-                        <span class="s2">&quot;b c n_patch_z n_patch_y n_patch_x -&gt; b c (n_patch_z patch_size_z) (n_patch_y patch_size_y) (n_patch_x patch_size_x)&quot;</span><span class="p">,</span>
-                        <span class="n">reduction</span><span class="o">=</span><span class="s2">&quot;repeat&quot;</span><span class="p">,</span>
-                        <span class="n">patch_size_z</span><span class="o">=</span><span class="n">patch_size</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span>
-                        <span class="n">patch_size_y</span><span class="o">=</span><span class="n">patch_size</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span>
-                        <span class="n">patch_size_x</span><span class="o">=</span><span class="n">patch_size</span><span class="p">[</span><span class="mi">2</span><span class="p">],</span>
-                    <span class="p">),</span>
-                <span class="p">]</span>
-            <span class="p">)</span>
+        <span class="k">if</span> <span class="n">np</span><span class="o">.</span><span class="n">any</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">q_stride</span> <span class="o">&gt;</span> <span class="mi">1</span><span class="p">):</span>
+            <span class="c1"># within a mask unit, tokens are spatially ordered</span>
+            <span class="c1"># perform spatial 2x2x2 max pooling over tokens</span>
+            <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">spatial_dims</span> <span class="o">==</span> <span class="mi">3</span><span class="p">:</span>
+                <span class="n">q</span> <span class="o">=</span> <span class="n">reduce</span><span class="p">(</span>
+                    <span class="n">q</span><span class="p">,</span>
+                    <span class="s2">&quot;batch num_mask_units num_heads (n_patches_z q_stride_z n_patches_y q_stride_y n_patches_x q_stride_x) c -&gt; batch num_mask_units num_heads (n_patches_z n_patches_y n_patches_x) c&quot;</span><span class="p">,</span>
+                    <span class="n">reduction</span><span class="o">=</span><span class="s2">&quot;max&quot;</span><span class="p">,</span>
+                    <span class="n">q_stride_z</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">q_stride</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span>
+                    <span class="n">q_stride_y</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">q_stride</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span>
+                    <span class="n">q_stride_x</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">q_stride</span><span class="p">[</span><span class="mi">2</span><span class="p">],</span>
+                    <span class="n">n_patches_z</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">pooled_patches_per_mask_unit</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span>
+                    <span class="n">n_patches_y</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">pooled_patches_per_mask_unit</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span>
+                    <span class="n">n_patches_x</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">pooled_patches_per_mask_unit</span><span class="p">[</span><span class="mi">2</span><span class="p">],</span>
+                <span class="p">)</span>
+            <span class="k">elif</span> <span class="bp">self</span><span class="o">.</span><span class="n">spatial_dims</span> <span class="o">==</span> <span class="mi">2</span><span class="p">:</span>
+                <span class="n">q</span> <span class="o">=</span> <span class="n">reduce</span><span class="p">(</span>
+                    <span class="n">q</span><span class="p">,</span>
+                    <span class="s2">&quot;batch num_mask_units num_heads (n_patches_y q_stride_y n_patches_x q_stride_x) c -&gt;batch num_mask_units num_heads (n_patches_y n_patches_x) c&quot;</span><span class="p">,</span>
+                    <span class="n">reduction</span><span class="o">=</span><span class="s2">&quot;max&quot;</span><span class="p">,</span>
+                    <span class="n">q_stride_y</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">q_stride</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span>
+                    <span class="n">q_stride_x</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">q_stride</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span>
+                    <span class="n">n_patches_y</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">pooled_patches_per_mask_unit</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span>
+                    <span class="n">n_patches_x</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">pooled_patches_per_mask_unit</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span>
+                <span class="p">)</span>
 
-        <span class="k">elif</span> <span class="n">spatial_dims</span> <span class="o">==</span> <span class="mi">2</span><span class="p">:</span>
-            <span class="bp">self</span><span class="o">.</span><span class="n">conv</span> <span class="o">=</span> <span class="n">nn</span><span class="o">.</span><span class="n">Conv2d</span><span class="p">(</span>
-                <span class="n">in_channels</span><span class="o">=</span><span class="n">input_channels</span><span class="p">,</span>
-                <span class="n">out_channels</span><span class="o">=</span><span class="n">emb_dim</span><span class="p">,</span>
-                <span class="n">kernel_size</span><span class="o">=</span><span class="n">weight_size</span><span class="p">,</span>
-                <span class="n">stride</span><span class="o">=</span><span class="n">patch_size</span><span class="p">,</span>
-                <span class="n">padding</span><span class="o">=</span><span class="n">context_pixels</span><span class="p">,</span>
-            <span class="p">)</span>
-            <span class="bp">self</span><span class="o">.</span><span class="n">img2token</span> <span class="o">=</span> <span class="n">Rearrange</span><span class="p">(</span><span class="s2">&quot;b c y x -&gt; (y x) b c&quot;</span><span class="p">)</span>
-            <span class="bp">self</span><span class="o">.</span><span class="n">patch2img</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Sequential</span><span class="p">(</span>
-                <span class="o">*</span><span class="p">[</span>
-                    <span class="n">Rearrange</span><span class="p">(</span>
-                        <span class="s2">&quot;(n_patch_y n_patch_x) b c -&gt;  b c n_patch_y n_patch_x&quot;</span><span class="p">,</span>
-                        <span class="n">n_patch_y</span><span class="o">=</span><span class="n">n_patches</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span>
-                        <span class="n">n_patch_x</span><span class="o">=</span><span class="n">n_patches</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span>
-                    <span class="p">),</span>
-                    <span class="n">Reduce</span><span class="p">(</span>
-                        <span class="s2">&quot;b c  n_patch_y n_patch_x -&gt; b c (n_patch_y patch_size_y) (n_patch_x patch_size_x)&quot;</span><span class="p">,</span>
-                        <span class="n">reduction</span><span class="o">=</span><span class="s2">&quot;repeat&quot;</span><span class="p">,</span>
-                        <span class="n">patch_size_y</span><span class="o">=</span><span class="n">patch_size</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span>
-                        <span class="n">patch_size_x</span><span class="o">=</span><span class="n">patch_size</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span>
-                    <span class="p">),</span>
-                <span class="p">]</span>
-            <span class="p">)</span>
-        <span class="bp">self</span><span class="o">.</span><span class="n">task_embedding</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">ParameterDict</span><span class="p">(</span>
-            <span class="p">{</span><span class="n">task</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Parameter</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="n">emb_dim</span><span class="p">))</span> <span class="k">for</span> <span class="n">task</span> <span class="ow">in</span> <span class="n">tasks</span><span class="p">}</span>
+        <span class="n">attn</span> <span class="o">=</span> <span class="n">F</span><span class="o">.</span><span class="n">scaled_dot_product_attention</span><span class="p">(</span>
+            <span class="n">q</span><span class="p">,</span>
+            <span class="n">k</span><span class="p">,</span>
+            <span class="n">v</span><span class="p">,</span>
+            <span class="n">dropout_p</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">attn_drop</span><span class="p">,</span>
+        <span class="p">)</span>
+        <span class="c1"># combine heads into single channel dimension</span>
+        <span class="n">x</span> <span class="o">=</span> <span class="n">rearrange</span><span class="p">(</span>
+            <span class="n">attn</span><span class="p">,</span> <span class="s2">&quot;b mask_units n_heads t c -&gt; b mask_units t (n_heads c)&quot;</span><span class="p">,</span> <span class="n">n_heads</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">num_heads</span>
         <span class="p">)</span>
-        <span class="bp">self</span><span class="o">.</span><span class="n">_init_weight</span><span class="p">()</span>
 
-    <span class="k">def</span> <span class="nf">_init_weight</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
-        <span class="k">for</span> <span class="n">task</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">task_embedding</span><span class="p">:</span>
-            <span class="n">trunc_normal_</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">task_embedding</span><span class="p">[</span><span class="n">task</span><span class="p">],</span> <span class="n">std</span><span class="o">=</span><span class="mf">0.02</span><span class="p">)</span>
+        <span class="n">x</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">proj</span><span class="p">(</span><span class="n">x</span><span class="p">)</span>
+        <span class="n">x</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">proj_drop</span><span class="p">(</span><span class="n">x</span><span class="p">)</span>
+        <span class="k">return</span> <span class="n">x</span></div></div>
 
-<div class="viewcode-block" id="Patchify.get_mask"><a class="viewcode-back" href="../../../../../cyto_dl.nn.vits.blocks.patchify.html#cyto_dl.nn.vits.blocks.patchify.Patchify.get_mask">[docs]</a>    <span class="k">def</span> <span class="nf">get_mask</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">img</span><span class="p">,</span> <span class="n">n_visible_patches</span><span class="p">,</span> <span class="n">num_patches</span><span class="p">):</span>
-        <span class="n">B</span> <span class="o">=</span> <span class="n">img</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
 
-        <span class="n">indexes</span> <span class="o">=</span> <span class="p">[</span><span class="n">random_indexes</span><span class="p">(</span><span class="n">num_patches</span><span class="p">,</span> <span class="n">img</span><span class="o">.</span><span class="n">device</span><span class="p">)</span> <span class="k">for</span> <span class="n">_</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">B</span><span class="p">)]</span>
-        <span class="c1"># forward indexes : index in image -&gt; shuffledpatch</span>
-        <span class="n">forward_indexes</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">stack</span><span class="p">([</span><span class="n">i</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="n">indexes</span><span class="p">],</span> <span class="n">axis</span><span class="o">=-</span><span class="mi">1</span><span class="p">)</span>
+<div class="viewcode-block" id="HieraBlock"><a class="viewcode-back" href="../../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html#cyto_dl.nn.vits.blocks.masked_unit_attention.HieraBlock">[docs]</a><span class="k">class</span> <span class="nc">HieraBlock</span><span class="p">(</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
+    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span>
+        <span class="bp">self</span><span class="p">,</span>
+        <span class="n">dim</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span>
+        <span class="n">dim_out</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span>
+        <span class="n">heads</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span>
+        <span class="n">spatial_dims</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">3</span><span class="p">,</span>
+        <span class="n">mlp_ratio</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">4.0</span><span class="p">,</span>
+        <span class="n">drop_path</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.0</span><span class="p">,</span>
+        <span class="n">norm_layer</span><span class="p">:</span> <span class="n">nn</span><span class="o">.</span><span class="n">Module</span> <span class="o">=</span> <span class="n">nn</span><span class="o">.</span><span class="n">LayerNorm</span><span class="p">,</span>
+        <span class="n">act_layer</span><span class="p">:</span> <span class="n">nn</span><span class="o">.</span><span class="n">Module</span> <span class="o">=</span> <span class="n">nn</span><span class="o">.</span><span class="n">GELU</span><span class="p">,</span>
+        <span class="n">q_stride</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="p">[</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">],</span>
+        <span class="n">patches_per_mask_unit</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="p">[</span><span class="mi">2</span><span class="p">,</span> <span class="mi">12</span><span class="p">,</span> <span class="mi">12</span><span class="p">],</span>
+    <span class="p">):</span>
+<span class="w">        </span><span class="sd">&quot;&quot;&quot;</span>
+<span class="sd">        Parameters</span>
+<span class="sd">        ----------</span>
+<span class="sd">        dim : int</span>
+<span class="sd">            Dimension of the input features.</span>
+<span class="sd">        dim_out : int</span>
+<span class="sd">            Dimension of the output features.</span>
+<span class="sd">        heads : int</span>
+<span class="sd">            Number of attention heads.</span>
+<span class="sd">        spatial_dims : int, optional</span>
+<span class="sd">            Number of spatial dimensions, by default 3.</span>
+<span class="sd">        mlp_ratio : float, optional</span>
+<span class="sd">            Ratio of MLP hidden dim to embedding dim, by default 4.0.</span>
+<span class="sd">        drop_path : float, optional</span>
+<span class="sd">            Dropout rate for the path, by default 0.0.</span>
+<span class="sd">        norm_layer : nn.Module, optional</span>
+<span class="sd">            Normalization layer, by default nn.LayerNorm.</span>
+<span class="sd">        act_layer : nn.Module, optional</span>
+<span class="sd">            Activation layer for the MLP, by default nn.GELU.</span>
+<span class="sd">        q_stride : List[int], optional</span>
+<span class="sd">            Stride for query, by default [1, 1, 1].</span>
+<span class="sd">        patches_per_mask_unit : List[int], optional</span>
+<span class="sd">            Number of patches per mask unit, by default [2, 12, 12].</span>
+<span class="sd">        &quot;&quot;&quot;</span>
+        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>
+        <span class="n">patches_per_mask_unit</span><span class="p">,</span> <span class="n">q_stride</span> <span class="o">=</span> <span class="n">match_tuple_dimensions</span><span class="p">(</span>
+            <span class="n">spatial_dims</span><span class="p">,</span> <span class="p">[</span><span class="n">patches_per_mask_unit</span><span class="p">,</span> <span class="n">q_stride</span><span class="p">]</span>
+        <span class="p">)</span>
 
-        <span class="c1"># backward indexes : shuffled patch -&gt; index in image</span>
-        <span class="n">backward_indexes</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">stack</span><span class="p">([</span><span class="n">i</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span> <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="n">indexes</span><span class="p">],</span> <span class="n">axis</span><span class="o">=-</span><span class="mi">1</span><span class="p">)</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">spatial_dims</span> <span class="o">=</span> <span class="n">spatial_dims</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">dim</span> <span class="o">=</span> <span class="n">dim</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">dim_out</span> <span class="o">=</span> <span class="n">dim_out</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">q_stride</span> <span class="o">=</span> <span class="n">q_stride</span>
 
-        <span class="n">mask</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="n">num_patches</span><span class="p">,</span> <span class="n">B</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="n">device</span><span class="o">=</span><span class="n">img</span><span class="o">.</span><span class="n">device</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">torch</span><span class="o">.</span><span class="n">bool</span><span class="p">)</span>
-        <span class="c1"># visible patches are first</span>
-        <span class="n">mask</span><span class="p">[:</span><span class="n">n_visible_patches</span><span class="p">]</span> <span class="o">=</span> <span class="kc">True</span>
-        <span class="n">mask</span> <span class="o">=</span> <span class="n">take_indexes</span><span class="p">(</span><span class="n">mask</span><span class="p">,</span> <span class="n">backward_indexes</span><span class="p">)</span>
-        <span class="n">mask</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">patch2img</span><span class="p">(</span><span class="n">mask</span><span class="p">)</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">norm1</span> <span class="o">=</span> <span class="n">norm_layer</span><span class="p">(</span><span class="n">dim</span><span class="p">)</span>
 
-        <span class="k">return</span> <span class="n">mask</span><span class="p">,</span> <span class="n">forward_indexes</span><span class="p">,</span> <span class="n">backward_indexes</span></div>
+        <span class="n">do_pool</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">any</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">(</span><span class="n">q_stride</span><span class="p">)</span> <span class="o">&gt;</span> <span class="mi">1</span><span class="p">)</span> <span class="ow">or</span> <span class="n">dim</span> <span class="o">!=</span> <span class="n">dim_out</span>
 
-<div class="viewcode-block" id="Patchify.forward"><a class="viewcode-back" href="../../../../../cyto_dl.nn.vits.blocks.patchify.html#cyto_dl.nn.vits.blocks.patchify.Patchify.forward">[docs]</a>    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">img</span><span class="p">,</span> <span class="n">mask_ratio</span><span class="p">,</span> <span class="n">task</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
-        <span class="c1"># generate mask</span>
-        <span class="n">num_patches</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">prod</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">n_patches</span><span class="p">)</span>
-        <span class="n">n_visible_patches</span> <span class="o">=</span> <span class="nb">int</span><span class="p">(</span><span class="n">num_patches</span> <span class="o">*</span> <span class="p">(</span><span class="mi">1</span> <span class="o">-</span> <span class="n">mask_ratio</span><span class="p">))</span>
-        <span class="n">mask</span><span class="p">,</span> <span class="n">forward_indexes</span><span class="p">,</span> <span class="n">backward_indexes</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">get_mask</span><span class="p">(</span>
-            <span class="n">img</span><span class="p">,</span> <span class="n">n_visible_patches</span><span class="p">,</span> <span class="n">num_patches</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">attn</span> <span class="o">=</span> <span class="n">MaskUnitAttention</span><span class="p">(</span>
+            <span class="n">dim</span><span class="p">,</span>
+            <span class="n">dim_out</span><span class="p">,</span>
+            <span class="n">spatial_dims</span><span class="o">=</span><span class="n">spatial_dims</span><span class="p">,</span>
+            <span class="n">num_heads</span><span class="o">=</span><span class="n">heads</span><span class="p">,</span>
+            <span class="n">q_stride</span><span class="o">=</span><span class="n">q_stride</span><span class="p">,</span>
+            <span class="n">patches_per_mask_unit</span><span class="o">=</span><span class="n">patches_per_mask_unit</span><span class="p">,</span>
         <span class="p">)</span>
-        <span class="c1"># generate patches</span>
-        <span class="n">tokens</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">conv</span><span class="p">(</span><span class="n">img</span> <span class="o">*</span> <span class="n">mask</span><span class="p">)</span>
-        <span class="n">tokens</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">img2token</span><span class="p">(</span><span class="n">tokens</span><span class="p">)</span>
-        <span class="c1"># add position embedding</span>
-        <span class="n">tokens</span> <span class="o">=</span> <span class="n">tokens</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">pos_embedding</span>
-        <span class="k">if</span> <span class="n">mask_ratio</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>
-            <span class="c1"># extract visible patches</span>
-            <span class="n">tokens</span> <span class="o">=</span> <span class="n">take_indexes</span><span class="p">(</span><span class="n">tokens</span><span class="p">,</span> <span class="n">forward_indexes</span><span class="p">)[:</span><span class="n">n_visible_patches</span><span class="p">]</span>
-
-        <span class="k">if</span> <span class="n">task</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">task_embedding</span><span class="p">:</span>
-            <span class="n">tokens</span> <span class="o">=</span> <span class="n">tokens</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">task_embedding</span><span class="p">[</span><span class="n">task</span><span class="p">]</span>
-
-        <span class="c1"># mask is used above to mask out patches, we need to invert it for loss calculation</span>
-        <span class="k">return</span> <span class="n">tokens</span><span class="p">,</span> <span class="o">~</span><span class="n">mask</span><span class="p">,</span> <span class="n">forward_indexes</span><span class="p">,</span> <span class="n">backward_indexes</span></div></div>
+
+        <span class="bp">self</span><span class="o">.</span><span class="n">norm2</span> <span class="o">=</span> <span class="n">norm_layer</span><span class="p">(</span><span class="n">dim_out</span><span class="p">)</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">mlp</span> <span class="o">=</span> <span class="n">Mlp</span><span class="p">(</span><span class="n">dim_out</span><span class="p">,</span> <span class="nb">int</span><span class="p">(</span><span class="n">dim_out</span> <span class="o">*</span> <span class="n">mlp_ratio</span><span class="p">),</span> <span class="n">act_layer</span><span class="o">=</span><span class="n">act_layer</span><span class="p">)</span>
+
+        <span class="bp">self</span><span class="o">.</span><span class="n">drop_path</span> <span class="o">=</span> <span class="n">DropPath</span><span class="p">(</span><span class="n">drop_path</span><span class="p">)</span> <span class="k">if</span> <span class="n">drop_path</span> <span class="o">&gt;</span> <span class="mi">0</span> <span class="k">else</span> <span class="n">nn</span><span class="o">.</span><span class="n">Identity</span><span class="p">()</span>
+
+        <span class="c1"># mean pooling by q stride within a mask unit</span>
+        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">spatial_dims</span> <span class="o">==</span> <span class="mi">3</span><span class="p">:</span>
+            <span class="n">skip_connection_pooling</span> <span class="o">=</span> <span class="n">Reduce</span><span class="p">(</span>
+                <span class="s2">&quot;b n (n_patches_z q_stride_z n_patches_y q_stride_y n_patches_x q_stride_x) c -&gt; b n (n_patches_z n_patches_y n_patches_x) c&quot;</span><span class="p">,</span>
+                <span class="n">reduction</span><span class="o">=</span><span class="s2">&quot;mean&quot;</span><span class="p">,</span>
+                <span class="n">q_stride_z</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">q_stride</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span>
+                <span class="n">q_stride_y</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">q_stride</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span>
+                <span class="n">q_stride_x</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">q_stride</span><span class="p">[</span><span class="mi">2</span><span class="p">],</span>
+                <span class="n">n_patches_z</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">attn</span><span class="o">.</span><span class="n">pooled_patches_per_mask_unit</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span>
+                <span class="n">n_patches_y</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">attn</span><span class="o">.</span><span class="n">pooled_patches_per_mask_unit</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span>
+                <span class="n">n_patches_x</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">attn</span><span class="o">.</span><span class="n">pooled_patches_per_mask_unit</span><span class="p">[</span><span class="mi">2</span><span class="p">],</span>
+            <span class="p">)</span>
+        <span class="k">elif</span> <span class="bp">self</span><span class="o">.</span><span class="n">spatial_dims</span> <span class="o">==</span> <span class="mi">2</span><span class="p">:</span>
+            <span class="n">skip_connection_pooling</span> <span class="o">=</span> <span class="n">Reduce</span><span class="p">(</span>
+                <span class="s2">&quot;b n (n_patches_y q_stride_y n_patches_x q_stride_x) c -&gt; b n (n_patches_y n_patches_x) c&quot;</span><span class="p">,</span>
+                <span class="n">reduction</span><span class="o">=</span><span class="s2">&quot;mean&quot;</span><span class="p">,</span>
+                <span class="n">q_stride_y</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">q_stride</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span>
+                <span class="n">q_stride_x</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">q_stride</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span>
+                <span class="n">n_patches_y</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">attn</span><span class="o">.</span><span class="n">pooled_patches_per_mask_unit</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span>
+                <span class="n">n_patches_x</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">attn</span><span class="o">.</span><span class="n">pooled_patches_per_mask_unit</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span>
+            <span class="p">)</span>
+
+        <span class="bp">self</span><span class="o">.</span><span class="n">proj</span> <span class="o">=</span> <span class="p">(</span>
+            <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Sequential</span><span class="p">(</span><span class="n">skip_connection_pooling</span><span class="p">,</span> <span class="n">nn</span><span class="o">.</span><span class="n">Linear</span><span class="p">(</span><span class="n">dim</span><span class="p">,</span> <span class="n">dim_out</span><span class="p">))</span>
+            <span class="k">if</span> <span class="n">do_pool</span>
+            <span class="k">else</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Identity</span><span class="p">()</span>
+        <span class="p">)</span>
+
+<div class="viewcode-block" id="HieraBlock.forward"><a class="viewcode-back" href="../../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html#cyto_dl.nn.vits.blocks.masked_unit_attention.HieraBlock.forward">[docs]</a>    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">:</span>
+<span class="w">        </span><span class="sd">&quot;&quot;&quot;</span>
+<span class="sd">        x: batch x mask units x tokens x emb_dim</span>
+<span class="sd">        &quot;&quot;&quot;</span>
+        <span class="c1"># Attention + Q Pooling</span>
+        <span class="n">x_norm</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">norm1</span><span class="p">(</span><span class="n">x</span><span class="p">)</span>
+
+        <span class="c1"># change dimension and subsample within mask unit for skip connection</span>
+        <span class="n">x</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">proj</span><span class="p">(</span><span class="n">x_norm</span><span class="p">)</span>
+
+        <span class="n">x</span> <span class="o">=</span> <span class="n">x</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">drop_path</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">attn</span><span class="p">(</span><span class="n">x_norm</span><span class="p">))</span>
+        <span class="c1"># MLP</span>
+        <span class="n">x</span> <span class="o">=</span> <span class="n">x</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">drop_path</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">mlp</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">norm2</span><span class="p">(</span><span class="n">x</span><span class="p">)))</span>
+        <span class="k">return</span> <span class="n">x</span></div></div>
 </pre></div>
         </article>
       </div>
diff --git a/_modules/cyto_dl/nn/vits/blocks/patchify/patchify.html b/_modules/cyto_dl/nn/vits/blocks/patchify/patchify.html
new file mode 100644
index 000000000..2fc7d8f94
--- /dev/null
+++ b/_modules/cyto_dl/nn/vits/blocks/patchify/patchify.html
@@ -0,0 +1,532 @@
+<!doctype html>
+<html class="no-js" lang="en">
+  <head><meta charset="utf-8"/>
+    <meta name="viewport" content="width=device-width,initial-scale=1"/>
+    <meta name="color-scheme" content="light dark"><link rel="index" title="Index" href="../../../../../../genindex.html" /><link rel="search" title="Search" href="../../../../../../search.html" />
+
+    <!-- Generated with Sphinx 6.1.3 and Furo 2023.03.23 -->
+        <title>cyto_dl.nn.vits.blocks.patchify.patchify - cyto-dl 0.4.0 documentation</title>
+      <link rel="stylesheet" type="text/css" href="../../../../../../_static/pygments.css" />
+    <link rel="stylesheet" type="text/css" href="../../../../../../_static/styles/furo.css?digest=fad236701ea90a88636c2a8c73b44ae642ed2a53" />
+    <link rel="stylesheet" type="text/css" href="../../../../../../_static/styles/furo-extensions.css?digest=30d1aed668e5c3a91c3e3bf6a60b675221979f0e" />
+    
+    
+
+
+<style>
+  body {
+    --color-code-background: #eeffcc;
+  --color-code-foreground: black;
+  
+  }
+  @media not print {
+    body[data-theme="dark"] {
+      --color-code-background: #202020;
+  --color-code-foreground: #d0d0d0;
+  
+    }
+    @media (prefers-color-scheme: dark) {
+      body:not([data-theme="light"]) {
+        --color-code-background: #202020;
+  --color-code-foreground: #d0d0d0;
+  
+      }
+    }
+  }
+</style></head>
+  <body>
+    
+    <script>
+      document.body.dataset.theme = localStorage.getItem("theme") || "auto";
+    </script>
+    
+
+<svg xmlns="http://www.w3.org/2000/svg" style="display: none;">
+  <symbol id="svg-toc" viewBox="0 0 24 24">
+    <title>Contents</title>
+    <svg stroke="currentColor" fill="currentColor" stroke-width="0" viewBox="0 0 1024 1024">
+      <path d="M408 442h480c4.4 0 8-3.6 8-8v-56c0-4.4-3.6-8-8-8H408c-4.4 0-8 3.6-8 8v56c0 4.4 3.6 8 8 8zm-8 204c0 4.4 3.6 8 8 8h480c4.4 0 8-3.6 8-8v-56c0-4.4-3.6-8-8-8H408c-4.4 0-8 3.6-8 8v56zm504-486H120c-4.4 0-8 3.6-8 8v56c0 4.4 3.6 8 8 8h784c4.4 0 8-3.6 8-8v-56c0-4.4-3.6-8-8-8zm0 632H120c-4.4 0-8 3.6-8 8v56c0 4.4 3.6 8 8 8h784c4.4 0 8-3.6 8-8v-56c0-4.4-3.6-8-8-8zM115.4 518.9L271.7 642c5.8 4.6 14.4.5 14.4-6.9V388.9c0-7.4-8.5-11.5-14.4-6.9L115.4 505.1a8.74 8.74 0 0 0 0 13.8z"/>
+    </svg>
+  </symbol>
+  <symbol id="svg-menu" viewBox="0 0 24 24">
+    <title>Menu</title>
+    <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
+      stroke-width="2" stroke-linecap="round" stroke-linejoin="round" class="feather-menu">
+      <line x1="3" y1="12" x2="21" y2="12"></line>
+      <line x1="3" y1="6" x2="21" y2="6"></line>
+      <line x1="3" y1="18" x2="21" y2="18"></line>
+    </svg>
+  </symbol>
+  <symbol id="svg-arrow-right" viewBox="0 0 24 24">
+    <title>Expand</title>
+    <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
+      stroke-width="2" stroke-linecap="round" stroke-linejoin="round" class="feather-chevron-right">
+      <polyline points="9 18 15 12 9 6"></polyline>
+    </svg>
+  </symbol>
+  <symbol id="svg-sun" viewBox="0 0 24 24">
+    <title>Light mode</title>
+    <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
+      stroke-width="1.5" stroke-linecap="round" stroke-linejoin="round" class="feather-sun">
+      <circle cx="12" cy="12" r="5"></circle>
+      <line x1="12" y1="1" x2="12" y2="3"></line>
+      <line x1="12" y1="21" x2="12" y2="23"></line>
+      <line x1="4.22" y1="4.22" x2="5.64" y2="5.64"></line>
+      <line x1="18.36" y1="18.36" x2="19.78" y2="19.78"></line>
+      <line x1="1" y1="12" x2="3" y2="12"></line>
+      <line x1="21" y1="12" x2="23" y2="12"></line>
+      <line x1="4.22" y1="19.78" x2="5.64" y2="18.36"></line>
+      <line x1="18.36" y1="5.64" x2="19.78" y2="4.22"></line>
+    </svg>
+  </symbol>
+  <symbol id="svg-moon" viewBox="0 0 24 24">
+    <title>Dark mode</title>
+    <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
+      stroke-width="1.5" stroke-linecap="round" stroke-linejoin="round" class="icon-tabler-moon">
+      <path stroke="none" d="M0 0h24v24H0z" fill="none" />
+      <path d="M12 3c.132 0 .263 0 .393 0a7.5 7.5 0 0 0 7.92 12.446a9 9 0 1 1 -8.313 -12.454z" />
+    </svg>
+  </symbol>
+  <symbol id="svg-sun-half" viewBox="0 0 24 24">
+    <title>Auto light/dark mode</title>
+    <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
+      stroke-width="1.5" stroke-linecap="round" stroke-linejoin="round" class="icon-tabler-shadow">
+      <path stroke="none" d="M0 0h24v24H0z" fill="none"/>
+      <circle cx="12" cy="12" r="9" />
+      <path d="M13 12h5" />
+      <path d="M13 15h4" />
+      <path d="M13 18h1" />
+      <path d="M13 9h4" />
+      <path d="M13 6h1" />
+    </svg>
+  </symbol>
+</svg>
+
+<input type="checkbox" class="sidebar-toggle" name="__navigation" id="__navigation">
+<input type="checkbox" class="sidebar-toggle" name="__toc" id="__toc">
+<label class="overlay sidebar-overlay" for="__navigation">
+  <div class="visually-hidden">Hide navigation sidebar</div>
+</label>
+<label class="overlay toc-overlay" for="__toc">
+  <div class="visually-hidden">Hide table of contents sidebar</div>
+</label>
+
+
+
+<div class="page">
+  <header class="mobile-header">
+    <div class="header-left">
+      <label class="nav-overlay-icon" for="__navigation">
+        <div class="visually-hidden">Toggle site navigation sidebar</div>
+        <i class="icon"><svg><use href="#svg-menu"></use></svg></i>
+      </label>
+    </div>
+    <div class="header-center">
+      <a href="../../../../../../index.html"><div class="brand">cyto-dl 0.4.0 documentation</div></a>
+    </div>
+    <div class="header-right">
+      <div class="theme-toggle-container theme-toggle-header">
+        <button class="theme-toggle">
+          <div class="visually-hidden">Toggle Light / Dark / Auto color theme</div>
+          <svg class="theme-icon-when-auto"><use href="#svg-sun-half"></use></svg>
+          <svg class="theme-icon-when-dark"><use href="#svg-moon"></use></svg>
+          <svg class="theme-icon-when-light"><use href="#svg-sun"></use></svg>
+        </button>
+      </div>
+      <label class="toc-overlay-icon toc-header-icon no-toc" for="__toc">
+        <div class="visually-hidden">Toggle table of contents sidebar</div>
+        <i class="icon"><svg><use href="#svg-toc"></use></svg></i>
+      </label>
+    </div>
+  </header>
+  <aside class="sidebar-drawer">
+    <div class="sidebar-container">
+      
+      <div class="sidebar-sticky"><a class="sidebar-brand" href="../../../../../../index.html">
+  
+  
+  <span class="sidebar-brand-text">cyto-dl 0.4.0 documentation</span>
+  
+</a><form class="sidebar-search-container" method="get" action="../../../../../../search.html" role="search">
+  <input class="sidebar-search" placeholder="Search" name="q" aria-label="Search">
+  <input type="hidden" name="check_keywords" value="yes">
+  <input type="hidden" name="area" value="default">
+</form>
+<div id="searchbox"></div><div class="sidebar-scroll"><div class="sidebar-tree">
+  <p class="caption" role="heading"><span class="caption-text">Contents:</span></p>
+<ul>
+<li class="toctree-l1"><a class="reference internal" href="../../../../../../index.html">Home</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../../../../../installation.html">Installation</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../../../../../using_examples.html">Modifying the Example Configs</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../../../../../api.html">Using the API</a></li>
+<li class="toctree-l1 has-children"><a class="reference internal" href="../../../../../../modules.html">Package modules</a><input class="toctree-checkbox" id="toctree-checkbox-1" name="toctree-checkbox-1" role="switch" type="checkbox"/><label for="toctree-checkbox-1"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l2 has-children"><a class="reference internal" href="../../../../../../cyto_dl.html">cyto_dl package</a><input class="toctree-checkbox" id="toctree-checkbox-2" name="toctree-checkbox-2" role="switch" type="checkbox"/><label for="toctree-checkbox-2"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../../cyto_dl.api.html">cyto_dl.api package</a><input class="toctree-checkbox" id="toctree-checkbox-3" name="toctree-checkbox-3" role="switch" type="checkbox"/><label for="toctree-checkbox-3"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.api.cyto_dl_model.html">cyto_dl.api.cyto_dl_model package</a><input class="toctree-checkbox" id="toctree-checkbox-4" name="toctree-checkbox-4" role="switch" type="checkbox"/><label for="toctree-checkbox-4"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.api.cyto_dl_model.cyto_dl_base_model.html">cyto_dl.api.cyto_dl_model.cyto_dl_base_model module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.api.cyto_dl_model.segmentation_plugin_model.html">cyto_dl.api.cyto_dl_model.segmentation_plugin_model module</a></li>
+</ul>
+</li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.api.data.html">cyto_dl.api.data module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.api.model.html">cyto_dl.api.model module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../../cyto_dl.callbacks.html">cyto_dl.callbacks package</a><input class="toctree-checkbox" id="toctree-checkbox-5" name="toctree-checkbox-5" role="switch" type="checkbox"/><label for="toctree-checkbox-5"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.callbacks.callback_utils.html">cyto_dl.callbacks.callback_utils module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.callbacks.csv_saver.html">cyto_dl.callbacks.csv_saver module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.callbacks.image_saver.html">cyto_dl.callbacks.image_saver module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.callbacks.latent_walk.html">cyto_dl.callbacks.latent_walk module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.callbacks.layer_freeze.html">cyto_dl.callbacks.layer_freeze module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.callbacks.model_utils.html">cyto_dl.callbacks.model_utils module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.callbacks.outlier_detection.html">cyto_dl.callbacks.outlier_detection module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../../cyto_dl.dataframe.html">cyto_dl.dataframe package</a><input class="toctree-checkbox" id="toctree-checkbox-6" name="toctree-checkbox-6" role="switch" type="checkbox"/><label for="toctree-checkbox-6"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.dataframe.transforms.html">cyto_dl.dataframe.transforms package</a><input class="toctree-checkbox" id="toctree-checkbox-7" name="toctree-checkbox-7" role="switch" type="checkbox"/><label for="toctree-checkbox-7"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.dataframe.transforms.filter.html">cyto_dl.dataframe.transforms.filter module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.dataframe.transforms.group_cols.html">cyto_dl.dataframe.transforms.group_cols module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.dataframe.transforms.misc.html">cyto_dl.dataframe.transforms.misc module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.dataframe.transforms.split.html">cyto_dl.dataframe.transforms.split module</a></li>
+</ul>
+</li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.dataframe.readers.html">cyto_dl.dataframe.readers module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../../cyto_dl.datamodules.html">cyto_dl.datamodules package</a><input class="toctree-checkbox" id="toctree-checkbox-8" name="toctree-checkbox-8" role="switch" type="checkbox"/><label for="toctree-checkbox-8"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.datamodules.dataframe.html">cyto_dl.datamodules.dataframe package</a><input class="toctree-checkbox" id="toctree-checkbox-9" name="toctree-checkbox-9" role="switch" type="checkbox"/><label for="toctree-checkbox-9"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.datamodules.dataframe.dataframe_datamodule.html">cyto_dl.datamodules.dataframe.dataframe_datamodule module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.datamodules.dataframe.grouped_dataframe_datamodule.html">cyto_dl.datamodules.dataframe.grouped_dataframe_datamodule module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.datamodules.dataframe.utils.html">cyto_dl.datamodules.dataframe.utils module</a></li>
+</ul>
+</li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.datamodules.array.html">cyto_dl.datamodules.array module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.datamodules.data_dict.html">cyto_dl.datamodules.data_dict module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.datamodules.folder.html">cyto_dl.datamodules.folder module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.datamodules.multidim_image.html">cyto_dl.datamodules.multidim_image module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.datamodules.smartcache.html">cyto_dl.datamodules.smartcache module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.datamodules.torchvision.html">cyto_dl.datamodules.torchvision module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../../cyto_dl.image.html">cyto_dl.image package</a><input class="toctree-checkbox" id="toctree-checkbox-10" name="toctree-checkbox-10" role="switch" type="checkbox"/><label for="toctree-checkbox-10"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.image.io.html">cyto_dl.image.io package</a><input class="toctree-checkbox" id="toctree-checkbox-11" name="toctree-checkbox-11" role="switch" type="checkbox"/><label for="toctree-checkbox-11"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.io.aicsimage_loader.html">cyto_dl.image.io.aicsimage_loader module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.io.monai_bio_reader.html">cyto_dl.image.io.monai_bio_reader module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.io.numpy_reader.html">cyto_dl.image.io.numpy_reader module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.io.ome_zarr_reader.html">cyto_dl.image.io.ome_zarr_reader module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.io.polygon_loader.html">cyto_dl.image.io.polygon_loader module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.io.skimage_reader.html">cyto_dl.image.io.skimage_reader module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.image.transforms.html">cyto_dl.image.transforms package</a><input class="toctree-checkbox" id="toctree-checkbox-12" name="toctree-checkbox-12" role="switch" type="checkbox"/><label for="toctree-checkbox-12"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.transforms.add_meta.html">cyto_dl.image.transforms.add_meta module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.transforms.bright_sampler.html">cyto_dl.image.transforms.bright_sampler module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.transforms.clip.html">cyto_dl.image.transforms.clip module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.transforms.contrastadjust.html">cyto_dl.image.transforms.contrastadjust module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.transforms.generate_jepa_masks.html">cyto_dl.image.transforms.generate_jepa_masks module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.transforms.merge.html">cyto_dl.image.transforms.merge module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.transforms.multiscale_cropper.html">cyto_dl.image.transforms.multiscale_cropper module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.transforms.pad.html">cyto_dl.image.transforms.pad module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.transforms.project.html">cyto_dl.image.transforms.project module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.transforms.rotation_mask_transform.html">cyto_dl.image.transforms.rotation_mask_transform module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.transforms.save.html">cyto_dl.image.transforms.save module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.transforms.so2_random_rotation.html">cyto_dl.image.transforms.so2_random_rotation module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.transforms.track_transforms.html">cyto_dl.image.transforms.track_transforms module</a></li>
+</ul>
+</li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../../cyto_dl.loggers.html">cyto_dl.loggers package</a><input class="toctree-checkbox" id="toctree-checkbox-13" name="toctree-checkbox-13" role="switch" type="checkbox"/><label for="toctree-checkbox-13"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.loggers.mlflow.html">cyto_dl.loggers.mlflow module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../../cyto_dl.models.html">cyto_dl.models package</a><input class="toctree-checkbox" id="toctree-checkbox-14" name="toctree-checkbox-14" role="switch" type="checkbox"/><label for="toctree-checkbox-14"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.models.classification.html">cyto_dl.models.classification package</a><input class="toctree-checkbox" id="toctree-checkbox-15" name="toctree-checkbox-15" role="switch" type="checkbox"/><label for="toctree-checkbox-15"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.classification.classification.html">cyto_dl.models.classification.classification module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.classification.timepoint_classification.html">cyto_dl.models.classification.timepoint_classification module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.models.contrastive.html">cyto_dl.models.contrastive package</a><input class="toctree-checkbox" id="toctree-checkbox-16" name="toctree-checkbox-16" role="switch" type="checkbox"/><label for="toctree-checkbox-16"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.contrastive.contrastive.html">cyto_dl.models.contrastive.contrastive module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.models.handlers.html">cyto_dl.models.handlers package</a><input class="toctree-checkbox" id="toctree-checkbox-17" name="toctree-checkbox-17" role="switch" type="checkbox"/><label for="toctree-checkbox-17"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.handlers.base_handler.html">cyto_dl.models.handlers.base_handler module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.handlers.image_handler.html">cyto_dl.models.handlers.image_handler module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.handlers.load_image_patch.html">cyto_dl.models.handlers.load_image_patch module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.models.im2im.html">cyto_dl.models.im2im package</a><input class="toctree-checkbox" id="toctree-checkbox-18" name="toctree-checkbox-18" role="switch" type="checkbox"/><label for="toctree-checkbox-18"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5 has-children"><a class="reference internal" href="../../../../../../cyto_dl.models.im2im.utils.html">cyto_dl.models.im2im.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-19" name="toctree-checkbox-19" role="switch" type="checkbox"/><label for="toctree-checkbox-19"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../../../cyto_dl.models.im2im.utils.postprocessing.html">cyto_dl.models.im2im.utils.postprocessing package</a><input class="toctree-checkbox" id="toctree-checkbox-20" name="toctree-checkbox-20" role="switch" type="checkbox"/><label for="toctree-checkbox-20"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../../../cyto_dl.models.im2im.utils.postprocessing.act_thresh_label.html">cyto_dl.models.im2im.utils.postprocessing.act_thresh_label module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../../../cyto_dl.models.im2im.utils.postprocessing.auto_thresh.html">cyto_dl.models.im2im.utils.postprocessing.auto_thresh module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../../../cyto_dl.models.im2im.utils.postprocessing.dict_to_im.html">cyto_dl.models.im2im.utils.postprocessing.dict_to_im module</a></li>
+</ul>
+</li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../../../cyto_dl.models.im2im.utils.instance_seg.html">cyto_dl.models.im2im.utils.instance_seg module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../../../cyto_dl.models.im2im.utils.noise_annealer.html">cyto_dl.models.im2im.utils.noise_annealer module</a></li>
+</ul>
+</li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.im2im.gan.html">cyto_dl.models.im2im.gan module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.im2im.multi_task.html">cyto_dl.models.im2im.multi_task module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.models.jepa.html">cyto_dl.models.jepa package</a><input class="toctree-checkbox" id="toctree-checkbox-21" name="toctree-checkbox-21" role="switch" type="checkbox"/><label for="toctree-checkbox-21"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.jepa.ijepa.html">cyto_dl.models.jepa.ijepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.jepa.iwm.html">cyto_dl.models.jepa.iwm module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.jepa.jepa_base.html">cyto_dl.models.jepa.jepa_base module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.models.utils.html">cyto_dl.models.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-22" name="toctree-checkbox-22" role="switch" type="checkbox"/><label for="toctree-checkbox-22"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.utils.mlflow.html">cyto_dl.models.utils.mlflow module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.utils.utils.html">cyto_dl.models.utils.utils module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.models.vae.html">cyto_dl.models.vae package</a><input class="toctree-checkbox" id="toctree-checkbox-23" name="toctree-checkbox-23" role="switch" type="checkbox"/><label for="toctree-checkbox-23"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5 has-children"><a class="reference internal" href="../../../../../../cyto_dl.models.vae.o2_spharm_vae.html">cyto_dl.models.vae.o2_spharm_vae package</a><input class="toctree-checkbox" id="toctree-checkbox-24" name="toctree-checkbox-24" role="switch" type="checkbox"/><label for="toctree-checkbox-24"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6"><a class="reference internal" href="../../../../../../cyto_dl.models.vae.o2_spharm_vae.o2_spharm_encoder.html">cyto_dl.models.vae.o2_spharm_vae.o2_spharm_encoder module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../../../cyto_dl.models.vae.o2_spharm_vae.o2_spharm_vae.html">cyto_dl.models.vae.o2_spharm_vae.o2_spharm_vae module</a></li>
+</ul>
+</li>
+<li class="toctree-l5 has-children"><a class="reference internal" href="../../../../../../cyto_dl.models.vae.priors.html">cyto_dl.models.vae.priors package</a><input class="toctree-checkbox" id="toctree-checkbox-25" name="toctree-checkbox-25" role="switch" type="checkbox"/><label for="toctree-checkbox-25"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6"><a class="reference internal" href="../../../../../../cyto_dl.models.vae.priors.abstract_prior.html">cyto_dl.models.vae.priors.abstract_prior module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../../../cyto_dl.models.vae.priors.gaussian.html">cyto_dl.models.vae.priors.gaussian module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../../../cyto_dl.models.vae.priors.identity_prior.html">cyto_dl.models.vae.priors.identity_prior module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../../../cyto_dl.models.vae.priors.joint_prior.html">cyto_dl.models.vae.priors.joint_prior module</a></li>
+</ul>
+</li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.vae.base_vae.html">cyto_dl.models.vae.base_vae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.vae.image_canon_vae.html">cyto_dl.models.vae.image_canon_vae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.vae.image_encoder.html">cyto_dl.models.vae.image_encoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.vae.image_vae.html">cyto_dl.models.vae.image_vae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.vae.latent_loss_vae.html">cyto_dl.models.vae.latent_loss_vae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.vae.point_cloud_vae.html">cyto_dl.models.vae.point_cloud_vae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.vae.tabular_vae.html">cyto_dl.models.vae.tabular_vae module</a></li>
+</ul>
+</li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.models.base_model.html">cyto_dl.models.base_model module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.models.basic_model.html">cyto_dl.models.basic_model module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../../cyto_dl.nn.html">cyto_dl.nn package</a><input class="toctree-checkbox" id="toctree-checkbox-26" name="toctree-checkbox-26" role="switch" type="checkbox"/><label for="toctree-checkbox-26"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.nn.discriminators.html">cyto_dl.nn.discriminators package</a><input class="toctree-checkbox" id="toctree-checkbox-27" name="toctree-checkbox-27" role="switch" type="checkbox"/><label for="toctree-checkbox-27"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.discriminators.multi_scale_discriminator.html">cyto_dl.nn.discriminators.multi_scale_discriminator module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.discriminators.n_layer_discriminator.html">cyto_dl.nn.discriminators.n_layer_discriminator module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.nn.head.html">cyto_dl.nn.head package</a><input class="toctree-checkbox" id="toctree-checkbox-28" name="toctree-checkbox-28" role="switch" type="checkbox"/><label for="toctree-checkbox-28"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.head.base_head.html">cyto_dl.nn.head.base_head module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.head.gan_head.html">cyto_dl.nn.head.gan_head module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.head.gan_head_superres.html">cyto_dl.nn.head.gan_head_superres module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.head.mae_head.html">cyto_dl.nn.head.mae_head module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.head.mask_head.html">cyto_dl.nn.head.mask_head module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.head.res_blocks_head.html">cyto_dl.nn.head.res_blocks_head module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.head.vic_reg.html">cyto_dl.nn.head.vic_reg module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.nn.losses.html">cyto_dl.nn.losses package</a><input class="toctree-checkbox" id="toctree-checkbox-29" name="toctree-checkbox-29" role="switch" type="checkbox"/><label for="toctree-checkbox-29"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.losses.adversarial_loss.html">cyto_dl.nn.losses.adversarial_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.losses.chamfer_loss.html">cyto_dl.nn.losses.chamfer_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.losses.continuous_bernoulli.html">cyto_dl.nn.losses.continuous_bernoulli module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.losses.cosine_loss.html">cyto_dl.nn.losses.cosine_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.losses.gan_loss.html">cyto_dl.nn.losses.gan_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.losses.gaussian_nll_loss.html">cyto_dl.nn.losses.gaussian_nll_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.losses.geomloss.html">cyto_dl.nn.losses.geomloss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.losses.loss_wrapper.html">cyto_dl.nn.losses.loss_wrapper module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.losses.spharm_loss.html">cyto_dl.nn.losses.spharm_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.losses.threshold_loss.html">cyto_dl.nn.losses.threshold_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.losses.vic_reg.html">cyto_dl.nn.losses.vic_reg module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.losses.weibull.html">cyto_dl.nn.losses.weibull module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.losses.weighted_mse_loss.html">cyto_dl.nn.losses.weighted_mse_loss module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.nn.point_cloud.html">cyto_dl.nn.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-30" name="toctree-checkbox-30" role="switch" type="checkbox"/><label for="toctree-checkbox-30"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.point_cloud.dgcnn.html">cyto_dl.nn.point_cloud.dgcnn module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.point_cloud.folding_net.html">cyto_dl.nn.point_cloud.folding_net module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.point_cloud.graph_functions.html">cyto_dl.nn.point_cloud.graph_functions module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.point_cloud.vnn.html">cyto_dl.nn.point_cloud.vnn module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5 has-children"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
+</ul>
+</li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
+</ul>
+</li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.nn.hr_skip.html">cyto_dl.nn.hr_skip module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.nn.mlp.html">cyto_dl.nn.mlp module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.nn.res_unit.html">cyto_dl.nn.res_unit module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.nn.spatial_transformer.html">cyto_dl.nn.spatial_transformer module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
+</ul>
+</li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
+</ul>
+</li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.utils.array.html">cyto_dl.utils.array module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.utils.checkpoint.html">cyto_dl.utils.checkpoint module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.utils.config.html">cyto_dl.utils.config module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.utils.download_test_data.html">cyto_dl.utils.download_test_data module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.utils.dummy_dataset.html">cyto_dl.utils.dummy_dataset module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.utils.pylogger.html">cyto_dl.utils.pylogger module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.utils.rich_utils.html">cyto_dl.utils.rich_utils module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.utils.rotation.html">cyto_dl.utils.rotation module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.utils.template_utils.html">cyto_dl.utils.template_utils module</a></li>
+</ul>
+</li>
+<li class="toctree-l3"><a class="reference internal" href="../../../../../../cyto_dl.compile.html">cyto_dl.compile module</a></li>
+<li class="toctree-l3"><a class="reference internal" href="../../../../../../cyto_dl.eval.html">cyto_dl.eval module</a></li>
+<li class="toctree-l3"><a class="reference internal" href="../../../../../../cyto_dl.train.html">cyto_dl.train module</a></li>
+</ul>
+</li>
+</ul>
+</li>
+</ul>
+
+</div>
+</div>
+
+      </div>
+      
+    </div>
+  </aside>
+  <div class="main">
+    <div class="content">
+      <div class="article-container">
+        <a href="#" class="back-to-top muted-link">
+          <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24">
+            <path d="M13 20h-2V8l-5.5 5.5-1.42-1.42L12 4.16l7.92 7.92-1.42 1.42L13 8v12z"></path>
+          </svg>
+          <span>Back to top</span>
+        </a>
+        <div class="content-icon-container">
+          <div class="theme-toggle-container theme-toggle-content">
+            <button class="theme-toggle">
+              <div class="visually-hidden">Toggle Light / Dark / Auto color theme</div>
+              <svg class="theme-icon-when-auto"><use href="#svg-sun-half"></use></svg>
+              <svg class="theme-icon-when-dark"><use href="#svg-moon"></use></svg>
+              <svg class="theme-icon-when-light"><use href="#svg-sun"></use></svg>
+            </button>
+          </div>
+          <label class="toc-overlay-icon toc-content-icon no-toc" for="__toc">
+            <div class="visually-hidden">Toggle table of contents sidebar</div>
+            <i class="icon"><svg><use href="#svg-toc"></use></svg></i>
+          </label>
+        </div>
+        <article role="main">
+          <h1>Source code for cyto_dl.nn.vits.blocks.patchify.patchify</h1><div class="highlight"><pre>
+<span></span><span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">List</span><span class="p">,</span> <span class="n">Optional</span>
+
+<span class="kn">import</span> <span class="nn">numpy</span> <span class="k">as</span> <span class="nn">np</span>
+<span class="kn">from</span> <span class="nn">einops.layers.torch</span> <span class="kn">import</span> <span class="n">Rearrange</span>
+
+<span class="kn">from</span> <span class="nn">cyto_dl.nn.vits.blocks.patchify.patchify_base</span> <span class="kn">import</span> <span class="n">PatchifyBase</span>
+<span class="kn">from</span> <span class="nn">cyto_dl.nn.vits.utils</span> <span class="kn">import</span> <span class="n">take_indexes</span>
+
+
+<div class="viewcode-block" id="Patchify"><a class="viewcode-back" href="../../../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html#cyto_dl.nn.vits.blocks.patchify.patchify.Patchify">[docs]</a><span class="k">class</span> <span class="nc">Patchify</span><span class="p">(</span><span class="n">PatchifyBase</span><span class="p">):</span>
+<span class="w">    </span><span class="sd">&quot;&quot;&quot;Class for converting images to a masked sequence of patches with positional embeddings.&quot;&quot;&quot;</span>
+
+    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span>
+        <span class="bp">self</span><span class="p">,</span>
+        <span class="n">patch_size</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span>
+        <span class="n">emb_dim</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span>
+        <span class="n">n_patches</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span>
+        <span class="n">spatial_dims</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">3</span><span class="p">,</span>
+        <span class="n">context_pixels</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">],</span>
+        <span class="n">input_channels</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">1</span><span class="p">,</span>
+        <span class="n">tasks</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">List</span><span class="p">[</span><span class="nb">str</span><span class="p">]]</span> <span class="o">=</span> <span class="p">[],</span>
+        <span class="n">learnable_pos_embedding</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">True</span><span class="p">,</span>
+    <span class="p">):</span>
+        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span>
+            <span class="n">patch_size</span><span class="o">=</span><span class="n">patch_size</span><span class="p">,</span>
+            <span class="n">emb_dim</span><span class="o">=</span><span class="n">emb_dim</span><span class="p">,</span>
+            <span class="n">n_patches</span><span class="o">=</span><span class="n">n_patches</span><span class="p">,</span>
+            <span class="n">spatial_dims</span><span class="o">=</span><span class="n">spatial_dims</span><span class="p">,</span>
+            <span class="n">context_pixels</span><span class="o">=</span><span class="n">context_pixels</span><span class="p">,</span>
+            <span class="n">input_channels</span><span class="o">=</span><span class="n">input_channels</span><span class="p">,</span>
+            <span class="n">tasks</span><span class="o">=</span><span class="n">tasks</span><span class="p">,</span>
+            <span class="n">learnable_pos_embedding</span><span class="o">=</span><span class="n">learnable_pos_embedding</span><span class="p">,</span>
+        <span class="p">)</span>
+
+    <span class="nd">@property</span>
+    <span class="k">def</span> <span class="nf">img2token</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
+        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">create_img2token</span><span class="p">()</span>
+
+<div class="viewcode-block" id="Patchify.get_mask_args"><a class="viewcode-back" href="../../../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html#cyto_dl.nn.vits.blocks.patchify.patchify.Patchify.get_mask_args">[docs]</a>    <span class="k">def</span> <span class="nf">get_mask_args</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">mask_ratio</span><span class="p">):</span>
+        <span class="n">num_patches</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">prod</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">n_patches</span><span class="p">)</span>
+        <span class="n">n_visible_patches</span> <span class="o">=</span> <span class="nb">int</span><span class="p">(</span><span class="n">num_patches</span> <span class="o">*</span> <span class="p">(</span><span class="mi">1</span> <span class="o">-</span> <span class="n">mask_ratio</span><span class="p">))</span>
+        <span class="k">return</span> <span class="n">n_visible_patches</span><span class="p">,</span> <span class="n">num_patches</span></div>
+
+<div class="viewcode-block" id="Patchify.create_img2token"><a class="viewcode-back" href="../../../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html#cyto_dl.nn.vits.blocks.patchify.patchify.Patchify.create_img2token">[docs]</a>    <span class="k">def</span> <span class="nf">create_img2token</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
+<span class="w">        </span><span class="sd">&quot;&quot;&quot;Rearranges the image tensor to a sequence of patches.&quot;&quot;&quot;</span>
+        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">spatial_dims</span> <span class="o">==</span> <span class="mi">3</span><span class="p">:</span>
+            <span class="k">return</span> <span class="n">Rearrange</span><span class="p">(</span><span class="s2">&quot;b c z y x -&gt; (z y x) b c&quot;</span><span class="p">)</span>
+        <span class="k">elif</span> <span class="bp">self</span><span class="o">.</span><span class="n">spatial_dims</span> <span class="o">==</span> <span class="mi">2</span><span class="p">:</span>
+            <span class="k">return</span> <span class="n">Rearrange</span><span class="p">(</span><span class="s2">&quot;b c y x -&gt; (y x) b c&quot;</span><span class="p">)</span></div>
+
+<div class="viewcode-block" id="Patchify.extract_visible_tokens"><a class="viewcode-back" href="../../../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html#cyto_dl.nn.vits.blocks.patchify.patchify.Patchify.extract_visible_tokens">[docs]</a>    <span class="k">def</span> <span class="nf">extract_visible_tokens</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">tokens</span><span class="p">,</span> <span class="n">forward_indexes</span><span class="p">,</span> <span class="n">n_visible_patches</span><span class="p">):</span>
+        <span class="k">return</span> <span class="n">take_indexes</span><span class="p">(</span><span class="n">tokens</span><span class="p">,</span> <span class="n">forward_indexes</span><span class="p">)[:</span><span class="n">n_visible_patches</span><span class="p">]</span></div></div>
+</pre></div>
+        </article>
+      </div>
+      <footer>
+        
+        <div class="related-pages">
+          
+          
+        </div>
+        <div class="bottom-of-page">
+          <div class="left-details">
+            <div class="copyright">
+                Copyright &#169; 2023
+            </div>
+            Made with <a href="https://www.sphinx-doc.org/">Sphinx</a> and <a class="muted-link" href="https://pradyunsg.me">@pradyunsg</a>'s
+            
+            <a href="https://github.com/pradyunsg/furo">Furo</a>
+            
+          </div>
+          <div class="right-details">
+            
+          </div>
+        </div>
+        
+      </footer>
+    </div>
+    <aside class="toc-drawer no-toc">
+      
+      
+      
+    </aside>
+  </div>
+</div><script data-url_root="../../../../../../" id="documentation_options" src="../../../../../../_static/documentation_options.js"></script>
+    <script src="../../../../../../_static/doctools.js"></script>
+    <script src="../../../../../../_static/sphinx_highlight.js"></script>
+    <script src="../../../../../../_static/scripts/furo.js"></script>
+    </body>
+</html>
\ No newline at end of file
diff --git a/_modules/cyto_dl/nn/vits/blocks/patchify/patchify_base.html b/_modules/cyto_dl/nn/vits/blocks/patchify/patchify_base.html
new file mode 100644
index 000000000..fb9faee2d
--- /dev/null
+++ b/_modules/cyto_dl/nn/vits/blocks/patchify/patchify_base.html
@@ -0,0 +1,667 @@
+<!doctype html>
+<html class="no-js" lang="en">
+  <head><meta charset="utf-8"/>
+    <meta name="viewport" content="width=device-width,initial-scale=1"/>
+    <meta name="color-scheme" content="light dark"><link rel="index" title="Index" href="../../../../../../genindex.html" /><link rel="search" title="Search" href="../../../../../../search.html" />
+
+    <!-- Generated with Sphinx 6.1.3 and Furo 2023.03.23 -->
+        <title>cyto_dl.nn.vits.blocks.patchify.patchify_base - cyto-dl 0.4.0 documentation</title>
+      <link rel="stylesheet" type="text/css" href="../../../../../../_static/pygments.css" />
+    <link rel="stylesheet" type="text/css" href="../../../../../../_static/styles/furo.css?digest=fad236701ea90a88636c2a8c73b44ae642ed2a53" />
+    <link rel="stylesheet" type="text/css" href="../../../../../../_static/styles/furo-extensions.css?digest=30d1aed668e5c3a91c3e3bf6a60b675221979f0e" />
+    
+    
+
+
+<style>
+  body {
+    --color-code-background: #eeffcc;
+  --color-code-foreground: black;
+  
+  }
+  @media not print {
+    body[data-theme="dark"] {
+      --color-code-background: #202020;
+  --color-code-foreground: #d0d0d0;
+  
+    }
+    @media (prefers-color-scheme: dark) {
+      body:not([data-theme="light"]) {
+        --color-code-background: #202020;
+  --color-code-foreground: #d0d0d0;
+  
+      }
+    }
+  }
+</style></head>
+  <body>
+    
+    <script>
+      document.body.dataset.theme = localStorage.getItem("theme") || "auto";
+    </script>
+    
+
+<svg xmlns="http://www.w3.org/2000/svg" style="display: none;">
+  <symbol id="svg-toc" viewBox="0 0 24 24">
+    <title>Contents</title>
+    <svg stroke="currentColor" fill="currentColor" stroke-width="0" viewBox="0 0 1024 1024">
+      <path d="M408 442h480c4.4 0 8-3.6 8-8v-56c0-4.4-3.6-8-8-8H408c-4.4 0-8 3.6-8 8v56c0 4.4 3.6 8 8 8zm-8 204c0 4.4 3.6 8 8 8h480c4.4 0 8-3.6 8-8v-56c0-4.4-3.6-8-8-8H408c-4.4 0-8 3.6-8 8v56zm504-486H120c-4.4 0-8 3.6-8 8v56c0 4.4 3.6 8 8 8h784c4.4 0 8-3.6 8-8v-56c0-4.4-3.6-8-8-8zm0 632H120c-4.4 0-8 3.6-8 8v56c0 4.4 3.6 8 8 8h784c4.4 0 8-3.6 8-8v-56c0-4.4-3.6-8-8-8zM115.4 518.9L271.7 642c5.8 4.6 14.4.5 14.4-6.9V388.9c0-7.4-8.5-11.5-14.4-6.9L115.4 505.1a8.74 8.74 0 0 0 0 13.8z"/>
+    </svg>
+  </symbol>
+  <symbol id="svg-menu" viewBox="0 0 24 24">
+    <title>Menu</title>
+    <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
+      stroke-width="2" stroke-linecap="round" stroke-linejoin="round" class="feather-menu">
+      <line x1="3" y1="12" x2="21" y2="12"></line>
+      <line x1="3" y1="6" x2="21" y2="6"></line>
+      <line x1="3" y1="18" x2="21" y2="18"></line>
+    </svg>
+  </symbol>
+  <symbol id="svg-arrow-right" viewBox="0 0 24 24">
+    <title>Expand</title>
+    <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
+      stroke-width="2" stroke-linecap="round" stroke-linejoin="round" class="feather-chevron-right">
+      <polyline points="9 18 15 12 9 6"></polyline>
+    </svg>
+  </symbol>
+  <symbol id="svg-sun" viewBox="0 0 24 24">
+    <title>Light mode</title>
+    <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
+      stroke-width="1.5" stroke-linecap="round" stroke-linejoin="round" class="feather-sun">
+      <circle cx="12" cy="12" r="5"></circle>
+      <line x1="12" y1="1" x2="12" y2="3"></line>
+      <line x1="12" y1="21" x2="12" y2="23"></line>
+      <line x1="4.22" y1="4.22" x2="5.64" y2="5.64"></line>
+      <line x1="18.36" y1="18.36" x2="19.78" y2="19.78"></line>
+      <line x1="1" y1="12" x2="3" y2="12"></line>
+      <line x1="21" y1="12" x2="23" y2="12"></line>
+      <line x1="4.22" y1="19.78" x2="5.64" y2="18.36"></line>
+      <line x1="18.36" y1="5.64" x2="19.78" y2="4.22"></line>
+    </svg>
+  </symbol>
+  <symbol id="svg-moon" viewBox="0 0 24 24">
+    <title>Dark mode</title>
+    <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
+      stroke-width="1.5" stroke-linecap="round" stroke-linejoin="round" class="icon-tabler-moon">
+      <path stroke="none" d="M0 0h24v24H0z" fill="none" />
+      <path d="M12 3c.132 0 .263 0 .393 0a7.5 7.5 0 0 0 7.92 12.446a9 9 0 1 1 -8.313 -12.454z" />
+    </svg>
+  </symbol>
+  <symbol id="svg-sun-half" viewBox="0 0 24 24">
+    <title>Auto light/dark mode</title>
+    <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
+      stroke-width="1.5" stroke-linecap="round" stroke-linejoin="round" class="icon-tabler-shadow">
+      <path stroke="none" d="M0 0h24v24H0z" fill="none"/>
+      <circle cx="12" cy="12" r="9" />
+      <path d="M13 12h5" />
+      <path d="M13 15h4" />
+      <path d="M13 18h1" />
+      <path d="M13 9h4" />
+      <path d="M13 6h1" />
+    </svg>
+  </symbol>
+</svg>
+
+<input type="checkbox" class="sidebar-toggle" name="__navigation" id="__navigation">
+<input type="checkbox" class="sidebar-toggle" name="__toc" id="__toc">
+<label class="overlay sidebar-overlay" for="__navigation">
+  <div class="visually-hidden">Hide navigation sidebar</div>
+</label>
+<label class="overlay toc-overlay" for="__toc">
+  <div class="visually-hidden">Hide table of contents sidebar</div>
+</label>
+
+
+
+<div class="page">
+  <header class="mobile-header">
+    <div class="header-left">
+      <label class="nav-overlay-icon" for="__navigation">
+        <div class="visually-hidden">Toggle site navigation sidebar</div>
+        <i class="icon"><svg><use href="#svg-menu"></use></svg></i>
+      </label>
+    </div>
+    <div class="header-center">
+      <a href="../../../../../../index.html"><div class="brand">cyto-dl 0.4.0 documentation</div></a>
+    </div>
+    <div class="header-right">
+      <div class="theme-toggle-container theme-toggle-header">
+        <button class="theme-toggle">
+          <div class="visually-hidden">Toggle Light / Dark / Auto color theme</div>
+          <svg class="theme-icon-when-auto"><use href="#svg-sun-half"></use></svg>
+          <svg class="theme-icon-when-dark"><use href="#svg-moon"></use></svg>
+          <svg class="theme-icon-when-light"><use href="#svg-sun"></use></svg>
+        </button>
+      </div>
+      <label class="toc-overlay-icon toc-header-icon no-toc" for="__toc">
+        <div class="visually-hidden">Toggle table of contents sidebar</div>
+        <i class="icon"><svg><use href="#svg-toc"></use></svg></i>
+      </label>
+    </div>
+  </header>
+  <aside class="sidebar-drawer">
+    <div class="sidebar-container">
+      
+      <div class="sidebar-sticky"><a class="sidebar-brand" href="../../../../../../index.html">
+  
+  
+  <span class="sidebar-brand-text">cyto-dl 0.4.0 documentation</span>
+  
+</a><form class="sidebar-search-container" method="get" action="../../../../../../search.html" role="search">
+  <input class="sidebar-search" placeholder="Search" name="q" aria-label="Search">
+  <input type="hidden" name="check_keywords" value="yes">
+  <input type="hidden" name="area" value="default">
+</form>
+<div id="searchbox"></div><div class="sidebar-scroll"><div class="sidebar-tree">
+  <p class="caption" role="heading"><span class="caption-text">Contents:</span></p>
+<ul>
+<li class="toctree-l1"><a class="reference internal" href="../../../../../../index.html">Home</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../../../../../installation.html">Installation</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../../../../../using_examples.html">Modifying the Example Configs</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../../../../../api.html">Using the API</a></li>
+<li class="toctree-l1 has-children"><a class="reference internal" href="../../../../../../modules.html">Package modules</a><input class="toctree-checkbox" id="toctree-checkbox-1" name="toctree-checkbox-1" role="switch" type="checkbox"/><label for="toctree-checkbox-1"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l2 has-children"><a class="reference internal" href="../../../../../../cyto_dl.html">cyto_dl package</a><input class="toctree-checkbox" id="toctree-checkbox-2" name="toctree-checkbox-2" role="switch" type="checkbox"/><label for="toctree-checkbox-2"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../../cyto_dl.api.html">cyto_dl.api package</a><input class="toctree-checkbox" id="toctree-checkbox-3" name="toctree-checkbox-3" role="switch" type="checkbox"/><label for="toctree-checkbox-3"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.api.cyto_dl_model.html">cyto_dl.api.cyto_dl_model package</a><input class="toctree-checkbox" id="toctree-checkbox-4" name="toctree-checkbox-4" role="switch" type="checkbox"/><label for="toctree-checkbox-4"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.api.cyto_dl_model.cyto_dl_base_model.html">cyto_dl.api.cyto_dl_model.cyto_dl_base_model module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.api.cyto_dl_model.segmentation_plugin_model.html">cyto_dl.api.cyto_dl_model.segmentation_plugin_model module</a></li>
+</ul>
+</li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.api.data.html">cyto_dl.api.data module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.api.model.html">cyto_dl.api.model module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../../cyto_dl.callbacks.html">cyto_dl.callbacks package</a><input class="toctree-checkbox" id="toctree-checkbox-5" name="toctree-checkbox-5" role="switch" type="checkbox"/><label for="toctree-checkbox-5"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.callbacks.callback_utils.html">cyto_dl.callbacks.callback_utils module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.callbacks.csv_saver.html">cyto_dl.callbacks.csv_saver module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.callbacks.image_saver.html">cyto_dl.callbacks.image_saver module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.callbacks.latent_walk.html">cyto_dl.callbacks.latent_walk module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.callbacks.layer_freeze.html">cyto_dl.callbacks.layer_freeze module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.callbacks.model_utils.html">cyto_dl.callbacks.model_utils module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.callbacks.outlier_detection.html">cyto_dl.callbacks.outlier_detection module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../../cyto_dl.dataframe.html">cyto_dl.dataframe package</a><input class="toctree-checkbox" id="toctree-checkbox-6" name="toctree-checkbox-6" role="switch" type="checkbox"/><label for="toctree-checkbox-6"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.dataframe.transforms.html">cyto_dl.dataframe.transforms package</a><input class="toctree-checkbox" id="toctree-checkbox-7" name="toctree-checkbox-7" role="switch" type="checkbox"/><label for="toctree-checkbox-7"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.dataframe.transforms.filter.html">cyto_dl.dataframe.transforms.filter module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.dataframe.transforms.group_cols.html">cyto_dl.dataframe.transforms.group_cols module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.dataframe.transforms.misc.html">cyto_dl.dataframe.transforms.misc module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.dataframe.transforms.split.html">cyto_dl.dataframe.transforms.split module</a></li>
+</ul>
+</li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.dataframe.readers.html">cyto_dl.dataframe.readers module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../../cyto_dl.datamodules.html">cyto_dl.datamodules package</a><input class="toctree-checkbox" id="toctree-checkbox-8" name="toctree-checkbox-8" role="switch" type="checkbox"/><label for="toctree-checkbox-8"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.datamodules.dataframe.html">cyto_dl.datamodules.dataframe package</a><input class="toctree-checkbox" id="toctree-checkbox-9" name="toctree-checkbox-9" role="switch" type="checkbox"/><label for="toctree-checkbox-9"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.datamodules.dataframe.dataframe_datamodule.html">cyto_dl.datamodules.dataframe.dataframe_datamodule module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.datamodules.dataframe.grouped_dataframe_datamodule.html">cyto_dl.datamodules.dataframe.grouped_dataframe_datamodule module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.datamodules.dataframe.utils.html">cyto_dl.datamodules.dataframe.utils module</a></li>
+</ul>
+</li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.datamodules.array.html">cyto_dl.datamodules.array module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.datamodules.data_dict.html">cyto_dl.datamodules.data_dict module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.datamodules.folder.html">cyto_dl.datamodules.folder module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.datamodules.multidim_image.html">cyto_dl.datamodules.multidim_image module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.datamodules.smartcache.html">cyto_dl.datamodules.smartcache module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.datamodules.torchvision.html">cyto_dl.datamodules.torchvision module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../../cyto_dl.image.html">cyto_dl.image package</a><input class="toctree-checkbox" id="toctree-checkbox-10" name="toctree-checkbox-10" role="switch" type="checkbox"/><label for="toctree-checkbox-10"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.image.io.html">cyto_dl.image.io package</a><input class="toctree-checkbox" id="toctree-checkbox-11" name="toctree-checkbox-11" role="switch" type="checkbox"/><label for="toctree-checkbox-11"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.io.aicsimage_loader.html">cyto_dl.image.io.aicsimage_loader module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.io.monai_bio_reader.html">cyto_dl.image.io.monai_bio_reader module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.io.numpy_reader.html">cyto_dl.image.io.numpy_reader module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.io.ome_zarr_reader.html">cyto_dl.image.io.ome_zarr_reader module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.io.polygon_loader.html">cyto_dl.image.io.polygon_loader module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.io.skimage_reader.html">cyto_dl.image.io.skimage_reader module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.image.transforms.html">cyto_dl.image.transforms package</a><input class="toctree-checkbox" id="toctree-checkbox-12" name="toctree-checkbox-12" role="switch" type="checkbox"/><label for="toctree-checkbox-12"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.transforms.add_meta.html">cyto_dl.image.transforms.add_meta module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.transforms.bright_sampler.html">cyto_dl.image.transforms.bright_sampler module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.transforms.clip.html">cyto_dl.image.transforms.clip module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.transforms.contrastadjust.html">cyto_dl.image.transforms.contrastadjust module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.transforms.generate_jepa_masks.html">cyto_dl.image.transforms.generate_jepa_masks module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.transforms.merge.html">cyto_dl.image.transforms.merge module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.transforms.multiscale_cropper.html">cyto_dl.image.transforms.multiscale_cropper module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.transforms.pad.html">cyto_dl.image.transforms.pad module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.transforms.project.html">cyto_dl.image.transforms.project module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.transforms.rotation_mask_transform.html">cyto_dl.image.transforms.rotation_mask_transform module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.transforms.save.html">cyto_dl.image.transforms.save module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.transforms.so2_random_rotation.html">cyto_dl.image.transforms.so2_random_rotation module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.transforms.track_transforms.html">cyto_dl.image.transforms.track_transforms module</a></li>
+</ul>
+</li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../../cyto_dl.loggers.html">cyto_dl.loggers package</a><input class="toctree-checkbox" id="toctree-checkbox-13" name="toctree-checkbox-13" role="switch" type="checkbox"/><label for="toctree-checkbox-13"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.loggers.mlflow.html">cyto_dl.loggers.mlflow module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../../cyto_dl.models.html">cyto_dl.models package</a><input class="toctree-checkbox" id="toctree-checkbox-14" name="toctree-checkbox-14" role="switch" type="checkbox"/><label for="toctree-checkbox-14"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.models.classification.html">cyto_dl.models.classification package</a><input class="toctree-checkbox" id="toctree-checkbox-15" name="toctree-checkbox-15" role="switch" type="checkbox"/><label for="toctree-checkbox-15"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.classification.classification.html">cyto_dl.models.classification.classification module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.classification.timepoint_classification.html">cyto_dl.models.classification.timepoint_classification module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.models.contrastive.html">cyto_dl.models.contrastive package</a><input class="toctree-checkbox" id="toctree-checkbox-16" name="toctree-checkbox-16" role="switch" type="checkbox"/><label for="toctree-checkbox-16"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.contrastive.contrastive.html">cyto_dl.models.contrastive.contrastive module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.models.handlers.html">cyto_dl.models.handlers package</a><input class="toctree-checkbox" id="toctree-checkbox-17" name="toctree-checkbox-17" role="switch" type="checkbox"/><label for="toctree-checkbox-17"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.handlers.base_handler.html">cyto_dl.models.handlers.base_handler module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.handlers.image_handler.html">cyto_dl.models.handlers.image_handler module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.handlers.load_image_patch.html">cyto_dl.models.handlers.load_image_patch module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.models.im2im.html">cyto_dl.models.im2im package</a><input class="toctree-checkbox" id="toctree-checkbox-18" name="toctree-checkbox-18" role="switch" type="checkbox"/><label for="toctree-checkbox-18"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5 has-children"><a class="reference internal" href="../../../../../../cyto_dl.models.im2im.utils.html">cyto_dl.models.im2im.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-19" name="toctree-checkbox-19" role="switch" type="checkbox"/><label for="toctree-checkbox-19"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../../../cyto_dl.models.im2im.utils.postprocessing.html">cyto_dl.models.im2im.utils.postprocessing package</a><input class="toctree-checkbox" id="toctree-checkbox-20" name="toctree-checkbox-20" role="switch" type="checkbox"/><label for="toctree-checkbox-20"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../../../cyto_dl.models.im2im.utils.postprocessing.act_thresh_label.html">cyto_dl.models.im2im.utils.postprocessing.act_thresh_label module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../../../cyto_dl.models.im2im.utils.postprocessing.auto_thresh.html">cyto_dl.models.im2im.utils.postprocessing.auto_thresh module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../../../cyto_dl.models.im2im.utils.postprocessing.dict_to_im.html">cyto_dl.models.im2im.utils.postprocessing.dict_to_im module</a></li>
+</ul>
+</li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../../../cyto_dl.models.im2im.utils.instance_seg.html">cyto_dl.models.im2im.utils.instance_seg module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../../../cyto_dl.models.im2im.utils.noise_annealer.html">cyto_dl.models.im2im.utils.noise_annealer module</a></li>
+</ul>
+</li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.im2im.gan.html">cyto_dl.models.im2im.gan module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.im2im.multi_task.html">cyto_dl.models.im2im.multi_task module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.models.jepa.html">cyto_dl.models.jepa package</a><input class="toctree-checkbox" id="toctree-checkbox-21" name="toctree-checkbox-21" role="switch" type="checkbox"/><label for="toctree-checkbox-21"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.jepa.ijepa.html">cyto_dl.models.jepa.ijepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.jepa.iwm.html">cyto_dl.models.jepa.iwm module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.jepa.jepa_base.html">cyto_dl.models.jepa.jepa_base module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.models.utils.html">cyto_dl.models.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-22" name="toctree-checkbox-22" role="switch" type="checkbox"/><label for="toctree-checkbox-22"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.utils.mlflow.html">cyto_dl.models.utils.mlflow module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.utils.utils.html">cyto_dl.models.utils.utils module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.models.vae.html">cyto_dl.models.vae package</a><input class="toctree-checkbox" id="toctree-checkbox-23" name="toctree-checkbox-23" role="switch" type="checkbox"/><label for="toctree-checkbox-23"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5 has-children"><a class="reference internal" href="../../../../../../cyto_dl.models.vae.o2_spharm_vae.html">cyto_dl.models.vae.o2_spharm_vae package</a><input class="toctree-checkbox" id="toctree-checkbox-24" name="toctree-checkbox-24" role="switch" type="checkbox"/><label for="toctree-checkbox-24"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6"><a class="reference internal" href="../../../../../../cyto_dl.models.vae.o2_spharm_vae.o2_spharm_encoder.html">cyto_dl.models.vae.o2_spharm_vae.o2_spharm_encoder module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../../../cyto_dl.models.vae.o2_spharm_vae.o2_spharm_vae.html">cyto_dl.models.vae.o2_spharm_vae.o2_spharm_vae module</a></li>
+</ul>
+</li>
+<li class="toctree-l5 has-children"><a class="reference internal" href="../../../../../../cyto_dl.models.vae.priors.html">cyto_dl.models.vae.priors package</a><input class="toctree-checkbox" id="toctree-checkbox-25" name="toctree-checkbox-25" role="switch" type="checkbox"/><label for="toctree-checkbox-25"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6"><a class="reference internal" href="../../../../../../cyto_dl.models.vae.priors.abstract_prior.html">cyto_dl.models.vae.priors.abstract_prior module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../../../cyto_dl.models.vae.priors.gaussian.html">cyto_dl.models.vae.priors.gaussian module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../../../cyto_dl.models.vae.priors.identity_prior.html">cyto_dl.models.vae.priors.identity_prior module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../../../cyto_dl.models.vae.priors.joint_prior.html">cyto_dl.models.vae.priors.joint_prior module</a></li>
+</ul>
+</li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.vae.base_vae.html">cyto_dl.models.vae.base_vae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.vae.image_canon_vae.html">cyto_dl.models.vae.image_canon_vae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.vae.image_encoder.html">cyto_dl.models.vae.image_encoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.vae.image_vae.html">cyto_dl.models.vae.image_vae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.vae.latent_loss_vae.html">cyto_dl.models.vae.latent_loss_vae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.vae.point_cloud_vae.html">cyto_dl.models.vae.point_cloud_vae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.vae.tabular_vae.html">cyto_dl.models.vae.tabular_vae module</a></li>
+</ul>
+</li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.models.base_model.html">cyto_dl.models.base_model module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.models.basic_model.html">cyto_dl.models.basic_model module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../../cyto_dl.nn.html">cyto_dl.nn package</a><input class="toctree-checkbox" id="toctree-checkbox-26" name="toctree-checkbox-26" role="switch" type="checkbox"/><label for="toctree-checkbox-26"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.nn.discriminators.html">cyto_dl.nn.discriminators package</a><input class="toctree-checkbox" id="toctree-checkbox-27" name="toctree-checkbox-27" role="switch" type="checkbox"/><label for="toctree-checkbox-27"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.discriminators.multi_scale_discriminator.html">cyto_dl.nn.discriminators.multi_scale_discriminator module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.discriminators.n_layer_discriminator.html">cyto_dl.nn.discriminators.n_layer_discriminator module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.nn.head.html">cyto_dl.nn.head package</a><input class="toctree-checkbox" id="toctree-checkbox-28" name="toctree-checkbox-28" role="switch" type="checkbox"/><label for="toctree-checkbox-28"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.head.base_head.html">cyto_dl.nn.head.base_head module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.head.gan_head.html">cyto_dl.nn.head.gan_head module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.head.gan_head_superres.html">cyto_dl.nn.head.gan_head_superres module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.head.mae_head.html">cyto_dl.nn.head.mae_head module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.head.mask_head.html">cyto_dl.nn.head.mask_head module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.head.res_blocks_head.html">cyto_dl.nn.head.res_blocks_head module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.head.vic_reg.html">cyto_dl.nn.head.vic_reg module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.nn.losses.html">cyto_dl.nn.losses package</a><input class="toctree-checkbox" id="toctree-checkbox-29" name="toctree-checkbox-29" role="switch" type="checkbox"/><label for="toctree-checkbox-29"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.losses.adversarial_loss.html">cyto_dl.nn.losses.adversarial_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.losses.chamfer_loss.html">cyto_dl.nn.losses.chamfer_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.losses.continuous_bernoulli.html">cyto_dl.nn.losses.continuous_bernoulli module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.losses.cosine_loss.html">cyto_dl.nn.losses.cosine_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.losses.gan_loss.html">cyto_dl.nn.losses.gan_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.losses.gaussian_nll_loss.html">cyto_dl.nn.losses.gaussian_nll_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.losses.geomloss.html">cyto_dl.nn.losses.geomloss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.losses.loss_wrapper.html">cyto_dl.nn.losses.loss_wrapper module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.losses.spharm_loss.html">cyto_dl.nn.losses.spharm_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.losses.threshold_loss.html">cyto_dl.nn.losses.threshold_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.losses.vic_reg.html">cyto_dl.nn.losses.vic_reg module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.losses.weibull.html">cyto_dl.nn.losses.weibull module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.losses.weighted_mse_loss.html">cyto_dl.nn.losses.weighted_mse_loss module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.nn.point_cloud.html">cyto_dl.nn.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-30" name="toctree-checkbox-30" role="switch" type="checkbox"/><label for="toctree-checkbox-30"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.point_cloud.dgcnn.html">cyto_dl.nn.point_cloud.dgcnn module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.point_cloud.folding_net.html">cyto_dl.nn.point_cloud.folding_net module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.point_cloud.graph_functions.html">cyto_dl.nn.point_cloud.graph_functions module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.point_cloud.vnn.html">cyto_dl.nn.point_cloud.vnn module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5 has-children"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
+</ul>
+</li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
+</ul>
+</li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.nn.hr_skip.html">cyto_dl.nn.hr_skip module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.nn.mlp.html">cyto_dl.nn.mlp module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.nn.res_unit.html">cyto_dl.nn.res_unit module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.nn.spatial_transformer.html">cyto_dl.nn.spatial_transformer module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
+</ul>
+</li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
+</ul>
+</li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.utils.array.html">cyto_dl.utils.array module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.utils.checkpoint.html">cyto_dl.utils.checkpoint module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.utils.config.html">cyto_dl.utils.config module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.utils.download_test_data.html">cyto_dl.utils.download_test_data module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.utils.dummy_dataset.html">cyto_dl.utils.dummy_dataset module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.utils.pylogger.html">cyto_dl.utils.pylogger module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.utils.rich_utils.html">cyto_dl.utils.rich_utils module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.utils.rotation.html">cyto_dl.utils.rotation module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.utils.template_utils.html">cyto_dl.utils.template_utils module</a></li>
+</ul>
+</li>
+<li class="toctree-l3"><a class="reference internal" href="../../../../../../cyto_dl.compile.html">cyto_dl.compile module</a></li>
+<li class="toctree-l3"><a class="reference internal" href="../../../../../../cyto_dl.eval.html">cyto_dl.eval module</a></li>
+<li class="toctree-l3"><a class="reference internal" href="../../../../../../cyto_dl.train.html">cyto_dl.train module</a></li>
+</ul>
+</li>
+</ul>
+</li>
+</ul>
+
+</div>
+</div>
+
+      </div>
+      
+    </div>
+  </aside>
+  <div class="main">
+    <div class="content">
+      <div class="article-container">
+        <a href="#" class="back-to-top muted-link">
+          <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24">
+            <path d="M13 20h-2V8l-5.5 5.5-1.42-1.42L12 4.16l7.92 7.92-1.42 1.42L13 8v12z"></path>
+          </svg>
+          <span>Back to top</span>
+        </a>
+        <div class="content-icon-container">
+          <div class="theme-toggle-container theme-toggle-content">
+            <button class="theme-toggle">
+              <div class="visually-hidden">Toggle Light / Dark / Auto color theme</div>
+              <svg class="theme-icon-when-auto"><use href="#svg-sun-half"></use></svg>
+              <svg class="theme-icon-when-dark"><use href="#svg-moon"></use></svg>
+              <svg class="theme-icon-when-light"><use href="#svg-sun"></use></svg>
+            </button>
+          </div>
+          <label class="toc-overlay-icon toc-content-icon no-toc" for="__toc">
+            <div class="visually-hidden">Toggle table of contents sidebar</div>
+            <i class="icon"><svg><use href="#svg-toc"></use></svg></i>
+          </label>
+        </div>
+        <article role="main">
+          <h1>Source code for cyto_dl.nn.vits.blocks.patchify.patchify_base</h1><div class="highlight"><pre>
+<span></span><span class="kn">from</span> <span class="nn">abc</span> <span class="kn">import</span> <span class="n">ABC</span><span class="p">,</span> <span class="n">abstractmethod</span>
+<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">List</span><span class="p">,</span> <span class="n">Optional</span>
+
+<span class="kn">import</span> <span class="nn">numpy</span> <span class="k">as</span> <span class="nn">np</span>
+<span class="kn">import</span> <span class="nn">torch</span>
+<span class="kn">import</span> <span class="nn">torch.nn</span> <span class="k">as</span> <span class="nn">nn</span>
+<span class="kn">from</span> <span class="nn">einops.layers.torch</span> <span class="kn">import</span> <span class="n">Rearrange</span><span class="p">,</span> <span class="n">Reduce</span>
+<span class="kn">from</span> <span class="nn">timm.models.layers</span> <span class="kn">import</span> <span class="n">trunc_normal_</span>
+
+<span class="kn">from</span> <span class="nn">cyto_dl.nn.vits.utils</span> <span class="kn">import</span> <span class="n">get_positional_embedding</span><span class="p">,</span> <span class="n">random_indexes</span><span class="p">,</span> <span class="n">take_indexes</span>
+
+
+<div class="viewcode-block" id="PatchifyBase"><a class="viewcode-back" href="../../../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html#cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase">[docs]</a><span class="k">class</span> <span class="nc">PatchifyBase</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">,</span> <span class="n">ABC</span><span class="p">):</span>
+<span class="w">    </span><span class="sd">&quot;&quot;&quot;Class for converting images to a masked sequence of patches with positional embeddings.&quot;&quot;&quot;</span>
+
+    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span>
+        <span class="bp">self</span><span class="p">,</span>
+        <span class="n">patch_size</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span>
+        <span class="n">emb_dim</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span>
+        <span class="n">n_patches</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span>
+        <span class="n">spatial_dims</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">3</span><span class="p">,</span>
+        <span class="n">context_pixels</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">],</span>
+        <span class="n">input_channels</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">1</span><span class="p">,</span>
+        <span class="n">tasks</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">List</span><span class="p">[</span><span class="nb">str</span><span class="p">]]</span> <span class="o">=</span> <span class="p">[],</span>
+        <span class="n">learnable_pos_embedding</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">True</span><span class="p">,</span>
+    <span class="p">):</span>
+<span class="w">        </span><span class="sd">&quot;&quot;&quot;</span>
+<span class="sd">        Parameters</span>
+<span class="sd">        ----------</span>
+<span class="sd">        patch_size: List[int]</span>
+<span class="sd">            Size of each patch in pix (ZYX order for 3D, YX order for 2D)</span>
+<span class="sd">        emb_dim: int</span>
+<span class="sd">            Dimension of encoder</span>
+<span class="sd">        n_patches: List[int]</span>
+<span class="sd">            Number of patches in each spatial dimension (ZYX order for 3D, YX order for 2D)</span>
+<span class="sd">        spatial_dims: int</span>
+<span class="sd">            Number of spatial dimensions</span>
+<span class="sd">        context_pixels: List[int]</span>
+<span class="sd">            Number of extra pixels around each patch to include in convolutional embedding to encoder dimension.</span>
+<span class="sd">        input_channels: int</span>
+<span class="sd">            Number of input channels</span>
+<span class="sd">        tasks: List[str]</span>
+<span class="sd">            List of tasks to encode</span>
+<span class="sd">        learnable_pos_embedding: bool</span>
+<span class="sd">            If True, learnable positional embeddings are used. If False, fixed sin/cos positional embeddings. Empirically, fixed positional embeddings work better for brightfield images.</span>
+<span class="sd">        &quot;&quot;&quot;</span>
+        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>
+
+        <span class="k">if</span> <span class="n">spatial_dims</span> <span class="ow">not</span> <span class="ow">in</span> <span class="p">(</span><span class="mi">2</span><span class="p">,</span> <span class="mi">3</span><span class="p">):</span>
+            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;Only 2D and 3D images are supported&quot;</span><span class="p">)</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">spatial_dims</span> <span class="o">=</span> <span class="n">spatial_dims</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">n_patches</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">asarray</span><span class="p">(</span><span class="n">n_patches</span><span class="p">)</span>
+
+        <span class="bp">self</span><span class="o">.</span><span class="n">pos_embedding</span> <span class="o">=</span> <span class="n">get_positional_embedding</span><span class="p">(</span>
+            <span class="n">n_patches</span><span class="p">,</span> <span class="n">emb_dim</span><span class="p">,</span> <span class="n">learnable</span><span class="o">=</span><span class="n">learnable_pos_embedding</span><span class="p">,</span> <span class="n">use_cls_token</span><span class="o">=</span><span class="kc">False</span>
+        <span class="p">)</span>
+
+        <span class="bp">self</span><span class="o">.</span><span class="n">patch2img</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">create_patch2img</span><span class="p">(</span><span class="n">n_patches</span><span class="p">,</span> <span class="n">patch_size</span><span class="p">)</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">conv</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">create_conv</span><span class="p">(</span><span class="n">input_channels</span><span class="p">,</span> <span class="n">emb_dim</span><span class="p">,</span> <span class="n">patch_size</span><span class="p">,</span> <span class="n">context_pixels</span><span class="p">)</span>
+
+        <span class="bp">self</span><span class="o">.</span><span class="n">task_embedding</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">ParameterDict</span><span class="p">(</span>
+            <span class="p">{</span><span class="n">task</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Parameter</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="n">emb_dim</span><span class="p">))</span> <span class="k">for</span> <span class="n">task</span> <span class="ow">in</span> <span class="n">tasks</span><span class="p">}</span>
+        <span class="p">)</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">_init_weight</span><span class="p">()</span>
+
+    <span class="k">def</span> <span class="nf">_init_weight</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
+        <span class="k">for</span> <span class="n">task</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">task_embedding</span><span class="p">:</span>
+            <span class="n">trunc_normal_</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">task_embedding</span><span class="p">[</span><span class="n">task</span><span class="p">],</span> <span class="n">std</span><span class="o">=</span><span class="mf">0.02</span><span class="p">)</span>
+
+    <span class="nd">@property</span>
+    <span class="nd">@abstractmethod</span>
+    <span class="k">def</span> <span class="nf">img2token</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
+        <span class="k">pass</span>
+
+<div class="viewcode-block" id="PatchifyBase.get_mask_args"><a class="viewcode-back" href="../../../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html#cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase.get_mask_args">[docs]</a>    <span class="nd">@abstractmethod</span>
+    <span class="k">def</span> <span class="nf">get_mask_args</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
+        <span class="k">pass</span></div>
+
+<div class="viewcode-block" id="PatchifyBase.extract_visible_tokens"><a class="viewcode-back" href="../../../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html#cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase.extract_visible_tokens">[docs]</a>    <span class="nd">@abstractmethod</span>
+    <span class="k">def</span> <span class="nf">extract_visible_tokens</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
+        <span class="k">pass</span></div>
+
+<div class="viewcode-block" id="PatchifyBase.create_conv"><a class="viewcode-back" href="../../../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html#cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase.create_conv">[docs]</a>    <span class="k">def</span> <span class="nf">create_conv</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">input_channels</span><span class="p">,</span> <span class="n">emb_dim</span><span class="p">,</span> <span class="n">patch_size</span><span class="p">,</span> <span class="n">context_pixels</span><span class="p">):</span>
+        <span class="n">context_pixels</span> <span class="o">=</span> <span class="n">context_pixels</span><span class="p">[:</span> <span class="bp">self</span><span class="o">.</span><span class="n">spatial_dims</span><span class="p">]</span>
+        <span class="n">weight_size</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">asarray</span><span class="p">(</span><span class="n">patch_size</span><span class="p">)</span> <span class="o">+</span> <span class="n">np</span><span class="o">.</span><span class="n">round</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">(</span><span class="n">context_pixels</span><span class="p">)</span> <span class="o">*</span> <span class="mi">2</span><span class="p">)</span><span class="o">.</span><span class="n">astype</span><span class="p">(</span><span class="nb">int</span><span class="p">)</span>
+
+        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">spatial_dims</span> <span class="o">==</span> <span class="mi">3</span><span class="p">:</span>
+            <span class="k">return</span> <span class="n">nn</span><span class="o">.</span><span class="n">Conv3d</span><span class="p">(</span>
+                <span class="n">in_channels</span><span class="o">=</span><span class="n">input_channels</span><span class="p">,</span>
+                <span class="n">out_channels</span><span class="o">=</span><span class="n">emb_dim</span><span class="p">,</span>
+                <span class="n">kernel_size</span><span class="o">=</span><span class="n">weight_size</span><span class="p">,</span>
+                <span class="n">stride</span><span class="o">=</span><span class="n">patch_size</span><span class="p">,</span>
+                <span class="n">padding</span><span class="o">=</span><span class="n">context_pixels</span><span class="p">,</span>
+            <span class="p">)</span>
+        <span class="k">elif</span> <span class="bp">self</span><span class="o">.</span><span class="n">spatial_dims</span> <span class="o">==</span> <span class="mi">2</span><span class="p">:</span>
+            <span class="k">return</span> <span class="n">nn</span><span class="o">.</span><span class="n">Conv2d</span><span class="p">(</span>
+                <span class="n">in_channels</span><span class="o">=</span><span class="n">input_channels</span><span class="p">,</span>
+                <span class="n">out_channels</span><span class="o">=</span><span class="n">emb_dim</span><span class="p">,</span>
+                <span class="n">kernel_size</span><span class="o">=</span><span class="n">weight_size</span><span class="p">,</span>
+                <span class="n">stride</span><span class="o">=</span><span class="n">patch_size</span><span class="p">,</span>
+                <span class="n">padding</span><span class="o">=</span><span class="n">context_pixels</span><span class="p">,</span>
+            <span class="p">)</span></div>
+
+<div class="viewcode-block" id="PatchifyBase.create_patch2img"><a class="viewcode-back" href="../../../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html#cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase.create_patch2img">[docs]</a>    <span class="k">def</span> <span class="nf">create_patch2img</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">n_patches</span><span class="p">,</span> <span class="n">patch_size</span><span class="p">):</span>
+<span class="w">        </span><span class="sd">&quot;&quot;&quot;Converts boolean array of whether to keep index of each patch to an image-shaped mask of</span>
+<span class="sd">        same size as input image.&quot;&quot;&quot;</span>
+        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">spatial_dims</span> <span class="o">==</span> <span class="mi">3</span><span class="p">:</span>
+            <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Sequential</span><span class="p">(</span>
+                <span class="o">*</span><span class="p">[</span>
+                    <span class="c1"># rearrange tokens to image</span>
+                    <span class="n">Rearrange</span><span class="p">(</span>
+                        <span class="s2">&quot;(n_patch_z n_patch_y n_patch_x) b c -&gt;  b c n_patch_z n_patch_y n_patch_x&quot;</span><span class="p">,</span>
+                        <span class="n">n_patch_z</span><span class="o">=</span><span class="n">n_patches</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span>
+                        <span class="n">n_patch_y</span><span class="o">=</span><span class="n">n_patches</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span>
+                        <span class="n">n_patch_x</span><span class="o">=</span><span class="n">n_patches</span><span class="p">[</span><span class="mi">2</span><span class="p">],</span>
+                    <span class="p">),</span>
+                    <span class="c1"># nearest neighbor resize image to match input image size</span>
+                    <span class="n">Reduce</span><span class="p">(</span>
+                        <span class="s2">&quot;b c n_patch_z n_patch_y n_patch_x -&gt; b c (n_patch_z patch_size_z) (n_patch_y patch_size_y) (n_patch_x patch_size_x)&quot;</span><span class="p">,</span>
+                        <span class="n">reduction</span><span class="o">=</span><span class="s2">&quot;repeat&quot;</span><span class="p">,</span>
+                        <span class="n">patch_size_z</span><span class="o">=</span><span class="n">patch_size</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span>
+                        <span class="n">patch_size_y</span><span class="o">=</span><span class="n">patch_size</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span>
+                        <span class="n">patch_size_x</span><span class="o">=</span><span class="n">patch_size</span><span class="p">[</span><span class="mi">2</span><span class="p">],</span>
+                    <span class="p">),</span>
+                <span class="p">]</span>
+            <span class="p">)</span>
+        <span class="k">elif</span> <span class="bp">self</span><span class="o">.</span><span class="n">spatial_dims</span> <span class="o">==</span> <span class="mi">2</span><span class="p">:</span>
+            <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Sequential</span><span class="p">(</span>
+                <span class="o">*</span><span class="p">[</span>
+                    <span class="n">Rearrange</span><span class="p">(</span>
+                        <span class="s2">&quot;(n_patch_y n_patch_x) b c -&gt;  b c n_patch_y n_patch_x&quot;</span><span class="p">,</span>
+                        <span class="n">n_patch_y</span><span class="o">=</span><span class="n">n_patches</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span>
+                        <span class="n">n_patch_x</span><span class="o">=</span><span class="n">n_patches</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span>
+                    <span class="p">),</span>
+                    <span class="n">Reduce</span><span class="p">(</span>
+                        <span class="s2">&quot;b c n_patch_y n_patch_x -&gt; b c (n_patch_y patch_size_y) (n_patch_x patch_size_x)&quot;</span><span class="p">,</span>
+                        <span class="n">reduction</span><span class="o">=</span><span class="s2">&quot;repeat&quot;</span><span class="p">,</span>
+                        <span class="n">patch_size_y</span><span class="o">=</span><span class="n">patch_size</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span>
+                        <span class="n">patch_size_x</span><span class="o">=</span><span class="n">patch_size</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span>
+                    <span class="p">),</span>
+                <span class="p">]</span>
+            <span class="p">)</span></div>
+
+<div class="viewcode-block" id="PatchifyBase.get_mask"><a class="viewcode-back" href="../../../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html#cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase.get_mask">[docs]</a>    <span class="k">def</span> <span class="nf">get_mask</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">img</span><span class="p">,</span> <span class="n">n_visible_patches</span><span class="p">,</span> <span class="n">num_patches</span><span class="p">):</span>
+        <span class="n">B</span> <span class="o">=</span> <span class="n">img</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
+
+        <span class="n">indexes</span> <span class="o">=</span> <span class="p">[</span><span class="n">random_indexes</span><span class="p">(</span><span class="n">num_patches</span><span class="p">,</span> <span class="n">img</span><span class="o">.</span><span class="n">device</span><span class="p">)</span> <span class="k">for</span> <span class="n">_</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">B</span><span class="p">)]</span>
+        <span class="c1"># forward indexes : index in image -&gt; shuffledpatch</span>
+        <span class="n">forward_indexes</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">stack</span><span class="p">([</span><span class="n">i</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="n">indexes</span><span class="p">],</span> <span class="n">axis</span><span class="o">=-</span><span class="mi">1</span><span class="p">)</span>
+
+        <span class="c1"># backward indexes : shuffled patch -&gt; index in image</span>
+        <span class="n">backward_indexes</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">stack</span><span class="p">([</span><span class="n">i</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span> <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="n">indexes</span><span class="p">],</span> <span class="n">axis</span><span class="o">=-</span><span class="mi">1</span><span class="p">)</span>
+
+        <span class="n">mask</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="n">num_patches</span><span class="p">,</span> <span class="n">B</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="n">device</span><span class="o">=</span><span class="n">img</span><span class="o">.</span><span class="n">device</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">torch</span><span class="o">.</span><span class="n">bool</span><span class="p">)</span>
+        <span class="c1"># visible patches are first</span>
+        <span class="n">mask</span><span class="p">[:</span><span class="n">n_visible_patches</span><span class="p">]</span> <span class="o">=</span> <span class="kc">True</span>
+        <span class="n">mask</span> <span class="o">=</span> <span class="n">take_indexes</span><span class="p">(</span><span class="n">mask</span><span class="p">,</span> <span class="n">backward_indexes</span><span class="p">)</span>
+        <span class="n">mask</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">patch2img</span><span class="p">(</span><span class="n">mask</span><span class="p">)</span>
+
+        <span class="k">return</span> <span class="n">mask</span><span class="p">,</span> <span class="n">forward_indexes</span><span class="p">,</span> <span class="n">backward_indexes</span></div>
+
+<div class="viewcode-block" id="PatchifyBase.forward"><a class="viewcode-back" href="../../../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html#cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase.forward">[docs]</a>    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">img</span><span class="p">,</span> <span class="n">mask_ratio</span><span class="p">,</span> <span class="n">task</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
+        <span class="c1"># generate mask</span>
+        <span class="n">mask</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">ones_like</span><span class="p">(</span><span class="n">img</span><span class="p">)</span><span class="o">.</span><span class="n">bool</span><span class="p">()</span>
+        <span class="n">forward_indexes</span><span class="p">,</span> <span class="n">backward_indexes</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="kc">None</span>
+        <span class="k">if</span> <span class="n">mask_ratio</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>
+            <span class="n">n_visible_patches</span><span class="p">,</span> <span class="n">num_patches</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">get_mask_args</span><span class="p">(</span><span class="n">mask_ratio</span><span class="p">)</span>
+            <span class="n">mask</span><span class="p">,</span> <span class="n">forward_indexes</span><span class="p">,</span> <span class="n">backward_indexes</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">get_mask</span><span class="p">(</span>
+                <span class="n">img</span><span class="p">,</span> <span class="n">n_visible_patches</span><span class="p">,</span> <span class="n">num_patches</span>
+            <span class="p">)</span>
+        <span class="c1"># generate patches</span>
+        <span class="n">tokens</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">conv</span><span class="p">(</span><span class="n">img</span> <span class="o">*</span> <span class="n">mask</span><span class="p">)</span>
+        <span class="n">tokens</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">img2token</span><span class="p">(</span><span class="n">tokens</span><span class="p">)</span>
+
+        <span class="c1"># add position embedding</span>
+        <span class="n">tokens</span> <span class="o">=</span> <span class="n">tokens</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">pos_embedding</span>
+
+        <span class="c1"># extract visible patches</span>
+        <span class="k">if</span> <span class="n">mask_ratio</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>
+            <span class="n">tokens</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">extract_visible_tokens</span><span class="p">(</span><span class="n">tokens</span><span class="p">,</span> <span class="n">forward_indexes</span><span class="p">,</span> <span class="n">n_visible_patches</span><span class="p">)</span>
+
+        <span class="c1"># add task embedding</span>
+        <span class="k">if</span> <span class="n">task</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">task_embedding</span><span class="p">:</span>
+            <span class="n">tokens</span> <span class="o">=</span> <span class="n">tokens</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">task_embedding</span><span class="p">[</span><span class="n">task</span><span class="p">]</span>
+
+        <span class="c1"># mask is used above to mask out patches, we need to invert it for loss calculation</span>
+        <span class="k">return</span> <span class="n">tokens</span><span class="p">,</span> <span class="o">~</span><span class="n">mask</span><span class="p">,</span> <span class="n">forward_indexes</span><span class="p">,</span> <span class="n">backward_indexes</span></div></div>
+</pre></div>
+        </article>
+      </div>
+      <footer>
+        
+        <div class="related-pages">
+          
+          
+        </div>
+        <div class="bottom-of-page">
+          <div class="left-details">
+            <div class="copyright">
+                Copyright &#169; 2023
+            </div>
+            Made with <a href="https://www.sphinx-doc.org/">Sphinx</a> and <a class="muted-link" href="https://pradyunsg.me">@pradyunsg</a>'s
+            
+            <a href="https://github.com/pradyunsg/furo">Furo</a>
+            
+          </div>
+          <div class="right-details">
+            
+          </div>
+        </div>
+        
+      </footer>
+    </div>
+    <aside class="toc-drawer no-toc">
+      
+      
+      
+    </aside>
+  </div>
+</div><script data-url_root="../../../../../../" id="documentation_options" src="../../../../../../_static/documentation_options.js"></script>
+    <script src="../../../../../../_static/doctools.js"></script>
+    <script src="../../../../../../_static/sphinx_highlight.js"></script>
+    <script src="../../../../../../_static/scripts/furo.js"></script>
+    </body>
+</html>
\ No newline at end of file
diff --git a/_modules/cyto_dl/nn/vits/blocks/patchify/patchify_hiera.html b/_modules/cyto_dl/nn/vits/blocks/patchify/patchify_hiera.html
new file mode 100644
index 000000000..8aabae9fa
--- /dev/null
+++ b/_modules/cyto_dl/nn/vits/blocks/patchify/patchify_hiera.html
@@ -0,0 +1,605 @@
+<!doctype html>
+<html class="no-js" lang="en">
+  <head><meta charset="utf-8"/>
+    <meta name="viewport" content="width=device-width,initial-scale=1"/>
+    <meta name="color-scheme" content="light dark"><link rel="index" title="Index" href="../../../../../../genindex.html" /><link rel="search" title="Search" href="../../../../../../search.html" />
+
+    <!-- Generated with Sphinx 6.1.3 and Furo 2023.03.23 -->
+        <title>cyto_dl.nn.vits.blocks.patchify.patchify_hiera - cyto-dl 0.4.0 documentation</title>
+      <link rel="stylesheet" type="text/css" href="../../../../../../_static/pygments.css" />
+    <link rel="stylesheet" type="text/css" href="../../../../../../_static/styles/furo.css?digest=fad236701ea90a88636c2a8c73b44ae642ed2a53" />
+    <link rel="stylesheet" type="text/css" href="../../../../../../_static/styles/furo-extensions.css?digest=30d1aed668e5c3a91c3e3bf6a60b675221979f0e" />
+    
+    
+
+
+<style>
+  body {
+    --color-code-background: #eeffcc;
+  --color-code-foreground: black;
+  
+  }
+  @media not print {
+    body[data-theme="dark"] {
+      --color-code-background: #202020;
+  --color-code-foreground: #d0d0d0;
+  
+    }
+    @media (prefers-color-scheme: dark) {
+      body:not([data-theme="light"]) {
+        --color-code-background: #202020;
+  --color-code-foreground: #d0d0d0;
+  
+      }
+    }
+  }
+</style></head>
+  <body>
+    
+    <script>
+      document.body.dataset.theme = localStorage.getItem("theme") || "auto";
+    </script>
+    
+
+<svg xmlns="http://www.w3.org/2000/svg" style="display: none;">
+  <symbol id="svg-toc" viewBox="0 0 24 24">
+    <title>Contents</title>
+    <svg stroke="currentColor" fill="currentColor" stroke-width="0" viewBox="0 0 1024 1024">
+      <path d="M408 442h480c4.4 0 8-3.6 8-8v-56c0-4.4-3.6-8-8-8H408c-4.4 0-8 3.6-8 8v56c0 4.4 3.6 8 8 8zm-8 204c0 4.4 3.6 8 8 8h480c4.4 0 8-3.6 8-8v-56c0-4.4-3.6-8-8-8H408c-4.4 0-8 3.6-8 8v56zm504-486H120c-4.4 0-8 3.6-8 8v56c0 4.4 3.6 8 8 8h784c4.4 0 8-3.6 8-8v-56c0-4.4-3.6-8-8-8zm0 632H120c-4.4 0-8 3.6-8 8v56c0 4.4 3.6 8 8 8h784c4.4 0 8-3.6 8-8v-56c0-4.4-3.6-8-8-8zM115.4 518.9L271.7 642c5.8 4.6 14.4.5 14.4-6.9V388.9c0-7.4-8.5-11.5-14.4-6.9L115.4 505.1a8.74 8.74 0 0 0 0 13.8z"/>
+    </svg>
+  </symbol>
+  <symbol id="svg-menu" viewBox="0 0 24 24">
+    <title>Menu</title>
+    <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
+      stroke-width="2" stroke-linecap="round" stroke-linejoin="round" class="feather-menu">
+      <line x1="3" y1="12" x2="21" y2="12"></line>
+      <line x1="3" y1="6" x2="21" y2="6"></line>
+      <line x1="3" y1="18" x2="21" y2="18"></line>
+    </svg>
+  </symbol>
+  <symbol id="svg-arrow-right" viewBox="0 0 24 24">
+    <title>Expand</title>
+    <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
+      stroke-width="2" stroke-linecap="round" stroke-linejoin="round" class="feather-chevron-right">
+      <polyline points="9 18 15 12 9 6"></polyline>
+    </svg>
+  </symbol>
+  <symbol id="svg-sun" viewBox="0 0 24 24">
+    <title>Light mode</title>
+    <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
+      stroke-width="1.5" stroke-linecap="round" stroke-linejoin="round" class="feather-sun">
+      <circle cx="12" cy="12" r="5"></circle>
+      <line x1="12" y1="1" x2="12" y2="3"></line>
+      <line x1="12" y1="21" x2="12" y2="23"></line>
+      <line x1="4.22" y1="4.22" x2="5.64" y2="5.64"></line>
+      <line x1="18.36" y1="18.36" x2="19.78" y2="19.78"></line>
+      <line x1="1" y1="12" x2="3" y2="12"></line>
+      <line x1="21" y1="12" x2="23" y2="12"></line>
+      <line x1="4.22" y1="19.78" x2="5.64" y2="18.36"></line>
+      <line x1="18.36" y1="5.64" x2="19.78" y2="4.22"></line>
+    </svg>
+  </symbol>
+  <symbol id="svg-moon" viewBox="0 0 24 24">
+    <title>Dark mode</title>
+    <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
+      stroke-width="1.5" stroke-linecap="round" stroke-linejoin="round" class="icon-tabler-moon">
+      <path stroke="none" d="M0 0h24v24H0z" fill="none" />
+      <path d="M12 3c.132 0 .263 0 .393 0a7.5 7.5 0 0 0 7.92 12.446a9 9 0 1 1 -8.313 -12.454z" />
+    </svg>
+  </symbol>
+  <symbol id="svg-sun-half" viewBox="0 0 24 24">
+    <title>Auto light/dark mode</title>
+    <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
+      stroke-width="1.5" stroke-linecap="round" stroke-linejoin="round" class="icon-tabler-shadow">
+      <path stroke="none" d="M0 0h24v24H0z" fill="none"/>
+      <circle cx="12" cy="12" r="9" />
+      <path d="M13 12h5" />
+      <path d="M13 15h4" />
+      <path d="M13 18h1" />
+      <path d="M13 9h4" />
+      <path d="M13 6h1" />
+    </svg>
+  </symbol>
+</svg>
+
+<input type="checkbox" class="sidebar-toggle" name="__navigation" id="__navigation">
+<input type="checkbox" class="sidebar-toggle" name="__toc" id="__toc">
+<label class="overlay sidebar-overlay" for="__navigation">
+  <div class="visually-hidden">Hide navigation sidebar</div>
+</label>
+<label class="overlay toc-overlay" for="__toc">
+  <div class="visually-hidden">Hide table of contents sidebar</div>
+</label>
+
+
+
+<div class="page">
+  <header class="mobile-header">
+    <div class="header-left">
+      <label class="nav-overlay-icon" for="__navigation">
+        <div class="visually-hidden">Toggle site navigation sidebar</div>
+        <i class="icon"><svg><use href="#svg-menu"></use></svg></i>
+      </label>
+    </div>
+    <div class="header-center">
+      <a href="../../../../../../index.html"><div class="brand">cyto-dl 0.4.0 documentation</div></a>
+    </div>
+    <div class="header-right">
+      <div class="theme-toggle-container theme-toggle-header">
+        <button class="theme-toggle">
+          <div class="visually-hidden">Toggle Light / Dark / Auto color theme</div>
+          <svg class="theme-icon-when-auto"><use href="#svg-sun-half"></use></svg>
+          <svg class="theme-icon-when-dark"><use href="#svg-moon"></use></svg>
+          <svg class="theme-icon-when-light"><use href="#svg-sun"></use></svg>
+        </button>
+      </div>
+      <label class="toc-overlay-icon toc-header-icon no-toc" for="__toc">
+        <div class="visually-hidden">Toggle table of contents sidebar</div>
+        <i class="icon"><svg><use href="#svg-toc"></use></svg></i>
+      </label>
+    </div>
+  </header>
+  <aside class="sidebar-drawer">
+    <div class="sidebar-container">
+      
+      <div class="sidebar-sticky"><a class="sidebar-brand" href="../../../../../../index.html">
+  
+  
+  <span class="sidebar-brand-text">cyto-dl 0.4.0 documentation</span>
+  
+</a><form class="sidebar-search-container" method="get" action="../../../../../../search.html" role="search">
+  <input class="sidebar-search" placeholder="Search" name="q" aria-label="Search">
+  <input type="hidden" name="check_keywords" value="yes">
+  <input type="hidden" name="area" value="default">
+</form>
+<div id="searchbox"></div><div class="sidebar-scroll"><div class="sidebar-tree">
+  <p class="caption" role="heading"><span class="caption-text">Contents:</span></p>
+<ul>
+<li class="toctree-l1"><a class="reference internal" href="../../../../../../index.html">Home</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../../../../../installation.html">Installation</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../../../../../using_examples.html">Modifying the Example Configs</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../../../../../api.html">Using the API</a></li>
+<li class="toctree-l1 has-children"><a class="reference internal" href="../../../../../../modules.html">Package modules</a><input class="toctree-checkbox" id="toctree-checkbox-1" name="toctree-checkbox-1" role="switch" type="checkbox"/><label for="toctree-checkbox-1"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l2 has-children"><a class="reference internal" href="../../../../../../cyto_dl.html">cyto_dl package</a><input class="toctree-checkbox" id="toctree-checkbox-2" name="toctree-checkbox-2" role="switch" type="checkbox"/><label for="toctree-checkbox-2"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../../cyto_dl.api.html">cyto_dl.api package</a><input class="toctree-checkbox" id="toctree-checkbox-3" name="toctree-checkbox-3" role="switch" type="checkbox"/><label for="toctree-checkbox-3"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.api.cyto_dl_model.html">cyto_dl.api.cyto_dl_model package</a><input class="toctree-checkbox" id="toctree-checkbox-4" name="toctree-checkbox-4" role="switch" type="checkbox"/><label for="toctree-checkbox-4"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.api.cyto_dl_model.cyto_dl_base_model.html">cyto_dl.api.cyto_dl_model.cyto_dl_base_model module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.api.cyto_dl_model.segmentation_plugin_model.html">cyto_dl.api.cyto_dl_model.segmentation_plugin_model module</a></li>
+</ul>
+</li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.api.data.html">cyto_dl.api.data module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.api.model.html">cyto_dl.api.model module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../../cyto_dl.callbacks.html">cyto_dl.callbacks package</a><input class="toctree-checkbox" id="toctree-checkbox-5" name="toctree-checkbox-5" role="switch" type="checkbox"/><label for="toctree-checkbox-5"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.callbacks.callback_utils.html">cyto_dl.callbacks.callback_utils module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.callbacks.csv_saver.html">cyto_dl.callbacks.csv_saver module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.callbacks.image_saver.html">cyto_dl.callbacks.image_saver module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.callbacks.latent_walk.html">cyto_dl.callbacks.latent_walk module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.callbacks.layer_freeze.html">cyto_dl.callbacks.layer_freeze module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.callbacks.model_utils.html">cyto_dl.callbacks.model_utils module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.callbacks.outlier_detection.html">cyto_dl.callbacks.outlier_detection module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../../cyto_dl.dataframe.html">cyto_dl.dataframe package</a><input class="toctree-checkbox" id="toctree-checkbox-6" name="toctree-checkbox-6" role="switch" type="checkbox"/><label for="toctree-checkbox-6"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.dataframe.transforms.html">cyto_dl.dataframe.transforms package</a><input class="toctree-checkbox" id="toctree-checkbox-7" name="toctree-checkbox-7" role="switch" type="checkbox"/><label for="toctree-checkbox-7"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.dataframe.transforms.filter.html">cyto_dl.dataframe.transforms.filter module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.dataframe.transforms.group_cols.html">cyto_dl.dataframe.transforms.group_cols module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.dataframe.transforms.misc.html">cyto_dl.dataframe.transforms.misc module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.dataframe.transforms.split.html">cyto_dl.dataframe.transforms.split module</a></li>
+</ul>
+</li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.dataframe.readers.html">cyto_dl.dataframe.readers module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../../cyto_dl.datamodules.html">cyto_dl.datamodules package</a><input class="toctree-checkbox" id="toctree-checkbox-8" name="toctree-checkbox-8" role="switch" type="checkbox"/><label for="toctree-checkbox-8"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.datamodules.dataframe.html">cyto_dl.datamodules.dataframe package</a><input class="toctree-checkbox" id="toctree-checkbox-9" name="toctree-checkbox-9" role="switch" type="checkbox"/><label for="toctree-checkbox-9"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.datamodules.dataframe.dataframe_datamodule.html">cyto_dl.datamodules.dataframe.dataframe_datamodule module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.datamodules.dataframe.grouped_dataframe_datamodule.html">cyto_dl.datamodules.dataframe.grouped_dataframe_datamodule module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.datamodules.dataframe.utils.html">cyto_dl.datamodules.dataframe.utils module</a></li>
+</ul>
+</li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.datamodules.array.html">cyto_dl.datamodules.array module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.datamodules.data_dict.html">cyto_dl.datamodules.data_dict module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.datamodules.folder.html">cyto_dl.datamodules.folder module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.datamodules.multidim_image.html">cyto_dl.datamodules.multidim_image module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.datamodules.smartcache.html">cyto_dl.datamodules.smartcache module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.datamodules.torchvision.html">cyto_dl.datamodules.torchvision module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../../cyto_dl.image.html">cyto_dl.image package</a><input class="toctree-checkbox" id="toctree-checkbox-10" name="toctree-checkbox-10" role="switch" type="checkbox"/><label for="toctree-checkbox-10"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.image.io.html">cyto_dl.image.io package</a><input class="toctree-checkbox" id="toctree-checkbox-11" name="toctree-checkbox-11" role="switch" type="checkbox"/><label for="toctree-checkbox-11"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.io.aicsimage_loader.html">cyto_dl.image.io.aicsimage_loader module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.io.monai_bio_reader.html">cyto_dl.image.io.monai_bio_reader module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.io.numpy_reader.html">cyto_dl.image.io.numpy_reader module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.io.ome_zarr_reader.html">cyto_dl.image.io.ome_zarr_reader module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.io.polygon_loader.html">cyto_dl.image.io.polygon_loader module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.io.skimage_reader.html">cyto_dl.image.io.skimage_reader module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.image.transforms.html">cyto_dl.image.transforms package</a><input class="toctree-checkbox" id="toctree-checkbox-12" name="toctree-checkbox-12" role="switch" type="checkbox"/><label for="toctree-checkbox-12"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.transforms.add_meta.html">cyto_dl.image.transforms.add_meta module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.transforms.bright_sampler.html">cyto_dl.image.transforms.bright_sampler module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.transforms.clip.html">cyto_dl.image.transforms.clip module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.transforms.contrastadjust.html">cyto_dl.image.transforms.contrastadjust module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.transforms.generate_jepa_masks.html">cyto_dl.image.transforms.generate_jepa_masks module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.transforms.merge.html">cyto_dl.image.transforms.merge module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.transforms.multiscale_cropper.html">cyto_dl.image.transforms.multiscale_cropper module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.transforms.pad.html">cyto_dl.image.transforms.pad module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.transforms.project.html">cyto_dl.image.transforms.project module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.transforms.rotation_mask_transform.html">cyto_dl.image.transforms.rotation_mask_transform module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.transforms.save.html">cyto_dl.image.transforms.save module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.transforms.so2_random_rotation.html">cyto_dl.image.transforms.so2_random_rotation module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.image.transforms.track_transforms.html">cyto_dl.image.transforms.track_transforms module</a></li>
+</ul>
+</li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../../cyto_dl.loggers.html">cyto_dl.loggers package</a><input class="toctree-checkbox" id="toctree-checkbox-13" name="toctree-checkbox-13" role="switch" type="checkbox"/><label for="toctree-checkbox-13"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.loggers.mlflow.html">cyto_dl.loggers.mlflow module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../../cyto_dl.models.html">cyto_dl.models package</a><input class="toctree-checkbox" id="toctree-checkbox-14" name="toctree-checkbox-14" role="switch" type="checkbox"/><label for="toctree-checkbox-14"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.models.classification.html">cyto_dl.models.classification package</a><input class="toctree-checkbox" id="toctree-checkbox-15" name="toctree-checkbox-15" role="switch" type="checkbox"/><label for="toctree-checkbox-15"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.classification.classification.html">cyto_dl.models.classification.classification module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.classification.timepoint_classification.html">cyto_dl.models.classification.timepoint_classification module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.models.contrastive.html">cyto_dl.models.contrastive package</a><input class="toctree-checkbox" id="toctree-checkbox-16" name="toctree-checkbox-16" role="switch" type="checkbox"/><label for="toctree-checkbox-16"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.contrastive.contrastive.html">cyto_dl.models.contrastive.contrastive module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.models.handlers.html">cyto_dl.models.handlers package</a><input class="toctree-checkbox" id="toctree-checkbox-17" name="toctree-checkbox-17" role="switch" type="checkbox"/><label for="toctree-checkbox-17"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.handlers.base_handler.html">cyto_dl.models.handlers.base_handler module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.handlers.image_handler.html">cyto_dl.models.handlers.image_handler module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.handlers.load_image_patch.html">cyto_dl.models.handlers.load_image_patch module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.models.im2im.html">cyto_dl.models.im2im package</a><input class="toctree-checkbox" id="toctree-checkbox-18" name="toctree-checkbox-18" role="switch" type="checkbox"/><label for="toctree-checkbox-18"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5 has-children"><a class="reference internal" href="../../../../../../cyto_dl.models.im2im.utils.html">cyto_dl.models.im2im.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-19" name="toctree-checkbox-19" role="switch" type="checkbox"/><label for="toctree-checkbox-19"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../../../cyto_dl.models.im2im.utils.postprocessing.html">cyto_dl.models.im2im.utils.postprocessing package</a><input class="toctree-checkbox" id="toctree-checkbox-20" name="toctree-checkbox-20" role="switch" type="checkbox"/><label for="toctree-checkbox-20"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../../../cyto_dl.models.im2im.utils.postprocessing.act_thresh_label.html">cyto_dl.models.im2im.utils.postprocessing.act_thresh_label module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../../../cyto_dl.models.im2im.utils.postprocessing.auto_thresh.html">cyto_dl.models.im2im.utils.postprocessing.auto_thresh module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../../../cyto_dl.models.im2im.utils.postprocessing.dict_to_im.html">cyto_dl.models.im2im.utils.postprocessing.dict_to_im module</a></li>
+</ul>
+</li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../../../cyto_dl.models.im2im.utils.instance_seg.html">cyto_dl.models.im2im.utils.instance_seg module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../../../cyto_dl.models.im2im.utils.noise_annealer.html">cyto_dl.models.im2im.utils.noise_annealer module</a></li>
+</ul>
+</li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.im2im.gan.html">cyto_dl.models.im2im.gan module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.im2im.multi_task.html">cyto_dl.models.im2im.multi_task module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.models.jepa.html">cyto_dl.models.jepa package</a><input class="toctree-checkbox" id="toctree-checkbox-21" name="toctree-checkbox-21" role="switch" type="checkbox"/><label for="toctree-checkbox-21"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.jepa.ijepa.html">cyto_dl.models.jepa.ijepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.jepa.iwm.html">cyto_dl.models.jepa.iwm module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.jepa.jepa_base.html">cyto_dl.models.jepa.jepa_base module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.models.utils.html">cyto_dl.models.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-22" name="toctree-checkbox-22" role="switch" type="checkbox"/><label for="toctree-checkbox-22"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.utils.mlflow.html">cyto_dl.models.utils.mlflow module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.utils.utils.html">cyto_dl.models.utils.utils module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.models.vae.html">cyto_dl.models.vae package</a><input class="toctree-checkbox" id="toctree-checkbox-23" name="toctree-checkbox-23" role="switch" type="checkbox"/><label for="toctree-checkbox-23"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5 has-children"><a class="reference internal" href="../../../../../../cyto_dl.models.vae.o2_spharm_vae.html">cyto_dl.models.vae.o2_spharm_vae package</a><input class="toctree-checkbox" id="toctree-checkbox-24" name="toctree-checkbox-24" role="switch" type="checkbox"/><label for="toctree-checkbox-24"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6"><a class="reference internal" href="../../../../../../cyto_dl.models.vae.o2_spharm_vae.o2_spharm_encoder.html">cyto_dl.models.vae.o2_spharm_vae.o2_spharm_encoder module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../../../cyto_dl.models.vae.o2_spharm_vae.o2_spharm_vae.html">cyto_dl.models.vae.o2_spharm_vae.o2_spharm_vae module</a></li>
+</ul>
+</li>
+<li class="toctree-l5 has-children"><a class="reference internal" href="../../../../../../cyto_dl.models.vae.priors.html">cyto_dl.models.vae.priors package</a><input class="toctree-checkbox" id="toctree-checkbox-25" name="toctree-checkbox-25" role="switch" type="checkbox"/><label for="toctree-checkbox-25"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6"><a class="reference internal" href="../../../../../../cyto_dl.models.vae.priors.abstract_prior.html">cyto_dl.models.vae.priors.abstract_prior module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../../../cyto_dl.models.vae.priors.gaussian.html">cyto_dl.models.vae.priors.gaussian module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../../../cyto_dl.models.vae.priors.identity_prior.html">cyto_dl.models.vae.priors.identity_prior module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../../../cyto_dl.models.vae.priors.joint_prior.html">cyto_dl.models.vae.priors.joint_prior module</a></li>
+</ul>
+</li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.vae.base_vae.html">cyto_dl.models.vae.base_vae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.vae.image_canon_vae.html">cyto_dl.models.vae.image_canon_vae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.vae.image_encoder.html">cyto_dl.models.vae.image_encoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.vae.image_vae.html">cyto_dl.models.vae.image_vae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.vae.latent_loss_vae.html">cyto_dl.models.vae.latent_loss_vae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.vae.point_cloud_vae.html">cyto_dl.models.vae.point_cloud_vae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.models.vae.tabular_vae.html">cyto_dl.models.vae.tabular_vae module</a></li>
+</ul>
+</li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.models.base_model.html">cyto_dl.models.base_model module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.models.basic_model.html">cyto_dl.models.basic_model module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../../cyto_dl.nn.html">cyto_dl.nn package</a><input class="toctree-checkbox" id="toctree-checkbox-26" name="toctree-checkbox-26" role="switch" type="checkbox"/><label for="toctree-checkbox-26"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.nn.discriminators.html">cyto_dl.nn.discriminators package</a><input class="toctree-checkbox" id="toctree-checkbox-27" name="toctree-checkbox-27" role="switch" type="checkbox"/><label for="toctree-checkbox-27"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.discriminators.multi_scale_discriminator.html">cyto_dl.nn.discriminators.multi_scale_discriminator module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.discriminators.n_layer_discriminator.html">cyto_dl.nn.discriminators.n_layer_discriminator module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.nn.head.html">cyto_dl.nn.head package</a><input class="toctree-checkbox" id="toctree-checkbox-28" name="toctree-checkbox-28" role="switch" type="checkbox"/><label for="toctree-checkbox-28"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.head.base_head.html">cyto_dl.nn.head.base_head module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.head.gan_head.html">cyto_dl.nn.head.gan_head module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.head.gan_head_superres.html">cyto_dl.nn.head.gan_head_superres module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.head.mae_head.html">cyto_dl.nn.head.mae_head module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.head.mask_head.html">cyto_dl.nn.head.mask_head module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.head.res_blocks_head.html">cyto_dl.nn.head.res_blocks_head module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.head.vic_reg.html">cyto_dl.nn.head.vic_reg module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.nn.losses.html">cyto_dl.nn.losses package</a><input class="toctree-checkbox" id="toctree-checkbox-29" name="toctree-checkbox-29" role="switch" type="checkbox"/><label for="toctree-checkbox-29"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.losses.adversarial_loss.html">cyto_dl.nn.losses.adversarial_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.losses.chamfer_loss.html">cyto_dl.nn.losses.chamfer_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.losses.continuous_bernoulli.html">cyto_dl.nn.losses.continuous_bernoulli module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.losses.cosine_loss.html">cyto_dl.nn.losses.cosine_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.losses.gan_loss.html">cyto_dl.nn.losses.gan_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.losses.gaussian_nll_loss.html">cyto_dl.nn.losses.gaussian_nll_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.losses.geomloss.html">cyto_dl.nn.losses.geomloss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.losses.loss_wrapper.html">cyto_dl.nn.losses.loss_wrapper module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.losses.spharm_loss.html">cyto_dl.nn.losses.spharm_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.losses.threshold_loss.html">cyto_dl.nn.losses.threshold_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.losses.vic_reg.html">cyto_dl.nn.losses.vic_reg module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.losses.weibull.html">cyto_dl.nn.losses.weibull module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.losses.weighted_mse_loss.html">cyto_dl.nn.losses.weighted_mse_loss module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.nn.point_cloud.html">cyto_dl.nn.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-30" name="toctree-checkbox-30" role="switch" type="checkbox"/><label for="toctree-checkbox-30"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.point_cloud.dgcnn.html">cyto_dl.nn.point_cloud.dgcnn module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.point_cloud.folding_net.html">cyto_dl.nn.point_cloud.folding_net module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.point_cloud.graph_functions.html">cyto_dl.nn.point_cloud.graph_functions module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.point_cloud.vnn.html">cyto_dl.nn.point_cloud.vnn module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5 has-children"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
+</ul>
+</li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
+</ul>
+</li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.nn.hr_skip.html">cyto_dl.nn.hr_skip module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.nn.mlp.html">cyto_dl.nn.mlp module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.nn.res_unit.html">cyto_dl.nn.res_unit module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.nn.spatial_transformer.html">cyto_dl.nn.spatial_transformer module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
+</ul>
+</li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
+</ul>
+</li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.utils.array.html">cyto_dl.utils.array module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.utils.checkpoint.html">cyto_dl.utils.checkpoint module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.utils.config.html">cyto_dl.utils.config module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.utils.download_test_data.html">cyto_dl.utils.download_test_data module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.utils.dummy_dataset.html">cyto_dl.utils.dummy_dataset module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.utils.pylogger.html">cyto_dl.utils.pylogger module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.utils.rich_utils.html">cyto_dl.utils.rich_utils module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.utils.rotation.html">cyto_dl.utils.rotation module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../../../cyto_dl.utils.template_utils.html">cyto_dl.utils.template_utils module</a></li>
+</ul>
+</li>
+<li class="toctree-l3"><a class="reference internal" href="../../../../../../cyto_dl.compile.html">cyto_dl.compile module</a></li>
+<li class="toctree-l3"><a class="reference internal" href="../../../../../../cyto_dl.eval.html">cyto_dl.eval module</a></li>
+<li class="toctree-l3"><a class="reference internal" href="../../../../../../cyto_dl.train.html">cyto_dl.train module</a></li>
+</ul>
+</li>
+</ul>
+</li>
+</ul>
+
+</div>
+</div>
+
+      </div>
+      
+    </div>
+  </aside>
+  <div class="main">
+    <div class="content">
+      <div class="article-container">
+        <a href="#" class="back-to-top muted-link">
+          <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24">
+            <path d="M13 20h-2V8l-5.5 5.5-1.42-1.42L12 4.16l7.92 7.92-1.42 1.42L13 8v12z"></path>
+          </svg>
+          <span>Back to top</span>
+        </a>
+        <div class="content-icon-container">
+          <div class="theme-toggle-container theme-toggle-content">
+            <button class="theme-toggle">
+              <div class="visually-hidden">Toggle Light / Dark / Auto color theme</div>
+              <svg class="theme-icon-when-auto"><use href="#svg-sun-half"></use></svg>
+              <svg class="theme-icon-when-dark"><use href="#svg-moon"></use></svg>
+              <svg class="theme-icon-when-light"><use href="#svg-sun"></use></svg>
+            </button>
+          </div>
+          <label class="toc-overlay-icon toc-content-icon no-toc" for="__toc">
+            <div class="visually-hidden">Toggle table of contents sidebar</div>
+            <i class="icon"><svg><use href="#svg-toc"></use></svg></i>
+          </label>
+        </div>
+        <article role="main">
+          <h1>Source code for cyto_dl.nn.vits.blocks.patchify.patchify_hiera</h1><div class="highlight"><pre>
+<span></span><span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">List</span><span class="p">,</span> <span class="n">Optional</span>
+
+<span class="kn">import</span> <span class="nn">numpy</span> <span class="k">as</span> <span class="nn">np</span>
+<span class="kn">import</span> <span class="nn">torch</span>
+<span class="kn">from</span> <span class="nn">einops</span> <span class="kn">import</span> <span class="n">rearrange</span><span class="p">,</span> <span class="n">repeat</span>
+<span class="kn">from</span> <span class="nn">einops.layers.torch</span> <span class="kn">import</span> <span class="n">Rearrange</span>
+<span class="kn">from</span> <span class="nn">timm.models.layers</span> <span class="kn">import</span> <span class="n">trunc_normal_</span>
+
+<span class="kn">from</span> <span class="nn">cyto_dl.nn.vits.blocks.patchify.patchify_base</span> <span class="kn">import</span> <span class="n">PatchifyBase</span>
+
+
+<div class="viewcode-block" id="take_indexes_mask"><a class="viewcode-back" href="../../../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html#cyto_dl.nn.vits.blocks.patchify.patchify_hiera.take_indexes_mask">[docs]</a><span class="k">def</span> <span class="nf">take_indexes_mask</span><span class="p">(</span><span class="n">sequences</span><span class="p">,</span> <span class="n">indexes</span><span class="p">):</span>
+<span class="w">    </span><span class="sd">&quot;&quot;&quot;</span>
+<span class="sd">    sequences: batch x mask units x patches x emb_dim</span>
+<span class="sd">    indexes: mask_units x batch</span>
+<span class="sd">    &quot;&quot;&quot;</span>
+    <span class="c1"># always gather across tokens dimension</span>
+    <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">gather</span><span class="p">(</span>
+        <span class="n">sequences</span><span class="p">,</span>
+        <span class="mi">1</span><span class="p">,</span>
+        <span class="n">repeat</span><span class="p">(</span>
+            <span class="n">indexes</span><span class="p">,</span>
+            <span class="s2">&quot;mu b -&gt; b mu p c&quot;</span><span class="p">,</span>
+            <span class="n">b</span><span class="o">=</span><span class="n">sequences</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span>
+            <span class="n">c</span><span class="o">=</span><span class="n">sequences</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">],</span>
+            <span class="n">mu</span><span class="o">=</span><span class="n">sequences</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span>
+            <span class="n">p</span><span class="o">=</span><span class="n">sequences</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">2</span><span class="p">],</span>
+        <span class="p">),</span>
+    <span class="p">)</span></div>
+
+
+<div class="viewcode-block" id="PatchifyHiera"><a class="viewcode-back" href="../../../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html#cyto_dl.nn.vits.blocks.patchify.patchify_hiera.PatchifyHiera">[docs]</a><span class="k">class</span> <span class="nc">PatchifyHiera</span><span class="p">(</span><span class="n">PatchifyBase</span><span class="p">):</span>
+<span class="w">    </span><span class="sd">&quot;&quot;&quot;Class for converting images to a sequence of patches with positional embeddings, masked at</span>
+<span class="sd">    the level of mask units (groups of patches specified by mask_units_per_dim).&quot;&quot;&quot;</span>
+
+    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span>
+        <span class="bp">self</span><span class="p">,</span>
+        <span class="n">patch_size</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span>
+        <span class="n">n_patches</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span>
+        <span class="n">emb_dim</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">64</span><span class="p">,</span>
+        <span class="n">spatial_dims</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">3</span><span class="p">,</span>
+        <span class="n">context_pixels</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">],</span>
+        <span class="n">input_channels</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">1</span><span class="p">,</span>
+        <span class="n">tasks</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">List</span><span class="p">[</span><span class="nb">str</span><span class="p">]]</span> <span class="o">=</span> <span class="p">[],</span>
+        <span class="n">mask_units_per_dim</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="p">[</span><span class="mi">8</span><span class="p">,</span> <span class="mi">8</span><span class="p">,</span> <span class="mi">8</span><span class="p">],</span>
+    <span class="p">):</span>
+<span class="w">        </span><span class="sd">&quot;&quot;&quot;</span>
+<span class="sd">        patch_size: List[int]</span>
+<span class="sd">            Size of each patch in pix (ZYX order for 3D, YX order for 2D)</span>
+<span class="sd">        n_patches: List[int]</span>
+<span class="sd">            Number of patches in each spatial dimension (ZYX order for 3D, YX order for 2D)</span>
+<span class="sd">        emb_dim: int</span>
+<span class="sd">            Dimension of encoder</span>
+<span class="sd">        spatial_dims: int</span>
+<span class="sd">            Number of spatial dimensions</span>
+<span class="sd">        context_pixels: List[int]</span>
+<span class="sd">            Number of extra pixels around each patch to include in convolutional embedding to encoder dimension.</span>
+<span class="sd">        input_channels: int</span>
+<span class="sd">            Number of input channels</span>
+<span class="sd">        tasks: List[str]</span>
+<span class="sd">            List of tasks to encode</span>
+<span class="sd">        mask_units_per_dim: List[int]</span>
+<span class="sd">            Number of mask units in each spatial dimension (ZYX order for 3D, YX order for 2D)</span>
+<span class="sd">        &quot;&quot;&quot;</span>
+        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span>
+            <span class="n">patch_size</span><span class="o">=</span><span class="n">patch_size</span><span class="p">,</span>
+            <span class="n">emb_dim</span><span class="o">=</span><span class="n">emb_dim</span><span class="p">,</span>
+            <span class="n">n_patches</span><span class="o">=</span><span class="n">n_patches</span><span class="p">,</span>
+            <span class="n">spatial_dims</span><span class="o">=</span><span class="n">spatial_dims</span><span class="p">,</span>
+            <span class="n">context_pixels</span><span class="o">=</span><span class="n">context_pixels</span><span class="p">,</span>
+            <span class="n">input_channels</span><span class="o">=</span><span class="n">input_channels</span><span class="p">,</span>
+            <span class="n">tasks</span><span class="o">=</span><span class="n">tasks</span><span class="p">,</span>
+            <span class="n">learnable_pos_embedding</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
+        <span class="p">)</span>
+
+        <span class="bp">self</span><span class="o">.</span><span class="n">total_n_mask_units</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">prod</span><span class="p">(</span><span class="n">mask_units_per_dim</span><span class="p">)</span>
+        <span class="c1"># mask_unit_size is the img shape / mask_units_per_dim, img_shape is size per patch * n_patches</span>
+        <span class="n">mask_unit_size_pix</span> <span class="o">=</span> <span class="p">(</span>
+            <span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">(</span><span class="n">patch_size</span><span class="p">)</span> <span class="o">*</span> <span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">(</span><span class="n">n_patches</span><span class="p">))</span> <span class="o">/</span> <span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">(</span><span class="n">mask_units_per_dim</span><span class="p">)</span>
+        <span class="p">)</span><span class="o">.</span><span class="n">astype</span><span class="p">(</span><span class="nb">int</span><span class="p">)</span>
+
+        <span class="n">patches_per_mask_unit</span> <span class="o">=</span> <span class="n">mask_unit_size_pix</span> <span class="o">//</span> <span class="n">patch_size</span>
+
+        <span class="c1"># rearrange patch embeddings to mask units</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">pos_embedding</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Parameter</span><span class="p">(</span>
+            <span class="n">rearrange</span><span class="p">(</span>
+                <span class="bp">self</span><span class="o">.</span><span class="n">pos_embedding</span><span class="p">,</span>
+                <span class="s2">&quot;(ppmu total_n_mu) 1 emb_dim -&gt; 1 total_n_mu ppmu emb_dim&quot;</span><span class="p">,</span>
+                <span class="n">total_n_mu</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">total_n_mask_units</span><span class="p">,</span>
+                <span class="n">ppmu</span><span class="o">=</span><span class="n">patches_per_mask_unit</span><span class="o">.</span><span class="n">prod</span><span class="p">(),</span>
+                <span class="n">emb_dim</span><span class="o">=</span><span class="n">emb_dim</span><span class="p">,</span>
+            <span class="p">)</span>
+        <span class="p">)</span>
+
+        <span class="bp">self</span><span class="o">.</span><span class="n">mask_units_per_dim</span> <span class="o">=</span> <span class="n">mask_units_per_dim</span>
+
+        <span class="bp">self</span><span class="o">.</span><span class="n">patch2img</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">create_patch2img</span><span class="p">(</span><span class="n">mask_units_per_dim</span><span class="p">,</span> <span class="n">mask_unit_size_pix</span><span class="p">)</span>
+
+    <span class="nd">@property</span>
+    <span class="k">def</span> <span class="nf">img2token</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
+        <span class="c1"># redefine this to work with mask units instead of patches</span>
+        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">create_img2token</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">mask_units_per_dim</span><span class="p">)</span>
+
+    <span class="c1"># in hiera, the masking is done at the mask unit, not the patch level</span>
+<div class="viewcode-block" id="PatchifyHiera.get_mask_args"><a class="viewcode-back" href="../../../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html#cyto_dl.nn.vits.blocks.patchify.patchify_hiera.PatchifyHiera.get_mask_args">[docs]</a>    <span class="k">def</span> <span class="nf">get_mask_args</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">mask_ratio</span><span class="p">):</span>
+        <span class="n">n_visible_patches</span> <span class="o">=</span> <span class="nb">int</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">total_n_mask_units</span> <span class="o">*</span> <span class="p">(</span><span class="mi">1</span> <span class="o">-</span> <span class="n">mask_ratio</span><span class="p">))</span>
+        <span class="k">return</span> <span class="n">n_visible_patches</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">total_n_mask_units</span></div>
+
+<div class="viewcode-block" id="PatchifyHiera.create_img2token"><a class="viewcode-back" href="../../../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html#cyto_dl.nn.vits.blocks.patchify.patchify_hiera.PatchifyHiera.create_img2token">[docs]</a>    <span class="k">def</span> <span class="nf">create_img2token</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">mask_units_per_dim</span><span class="p">):</span>
+        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">spatial_dims</span> <span class="o">==</span> <span class="mi">3</span><span class="p">:</span>
+            <span class="k">return</span> <span class="n">Rearrange</span><span class="p">(</span>
+                <span class="s2">&quot;b c (n_mu_z z) (n_mu_y y) (n_mu_x x) -&gt; b (n_mu_z n_mu_y n_mu_x) (z y x) c &quot;</span><span class="p">,</span>
+                <span class="n">n_mu_z</span><span class="o">=</span><span class="n">mask_units_per_dim</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span>
+                <span class="n">n_mu_y</span><span class="o">=</span><span class="n">mask_units_per_dim</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span>
+                <span class="n">n_mu_x</span><span class="o">=</span><span class="n">mask_units_per_dim</span><span class="p">[</span><span class="mi">2</span><span class="p">],</span>
+            <span class="p">)</span>
+        <span class="k">elif</span> <span class="bp">self</span><span class="o">.</span><span class="n">spatial_dims</span> <span class="o">==</span> <span class="mi">2</span><span class="p">:</span>
+            <span class="k">return</span> <span class="n">Rearrange</span><span class="p">(</span>
+                <span class="s2">&quot;b c  (n_mu_y y) (n_mu_x x) -&gt; b (n_mu_y n_mu_x) (y x) c &quot;</span><span class="p">,</span>
+                <span class="n">n_mu_y</span><span class="o">=</span><span class="n">mask_units_per_dim</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span>
+                <span class="n">n_mu_x</span><span class="o">=</span><span class="n">mask_units_per_dim</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span>
+            <span class="p">)</span></div>
+
+<div class="viewcode-block" id="PatchifyHiera.extract_visible_tokens"><a class="viewcode-back" href="../../../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html#cyto_dl.nn.vits.blocks.patchify.patchify_hiera.PatchifyHiera.extract_visible_tokens">[docs]</a>    <span class="k">def</span> <span class="nf">extract_visible_tokens</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">tokens</span><span class="p">,</span> <span class="n">forward_indexes</span><span class="p">,</span> <span class="n">n_visible_patches</span><span class="p">):</span>
+        <span class="k">return</span> <span class="n">take_indexes_mask</span><span class="p">(</span><span class="n">tokens</span><span class="p">,</span> <span class="n">forward_indexes</span><span class="p">)[:,</span> <span class="p">:</span><span class="n">n_visible_patches</span><span class="p">]</span></div></div>
+</pre></div>
+        </article>
+      </div>
+      <footer>
+        
+        <div class="related-pages">
+          
+          
+        </div>
+        <div class="bottom-of-page">
+          <div class="left-details">
+            <div class="copyright">
+                Copyright &#169; 2023
+            </div>
+            Made with <a href="https://www.sphinx-doc.org/">Sphinx</a> and <a class="muted-link" href="https://pradyunsg.me">@pradyunsg</a>'s
+            
+            <a href="https://github.com/pradyunsg/furo">Furo</a>
+            
+          </div>
+          <div class="right-details">
+            
+          </div>
+        </div>
+        
+      </footer>
+    </div>
+    <aside class="toc-drawer no-toc">
+      
+      
+      
+    </aside>
+  </div>
+</div><script data-url_root="../../../../../../" id="documentation_options" src="../../../../../../_static/documentation_options.js"></script>
+    <script src="../../../../../../_static/doctools.js"></script>
+    <script src="../../../../../../_static/sphinx_highlight.js"></script>
+    <script src="../../../../../../_static/scripts/furo.js"></script>
+    </body>
+</html>
\ No newline at end of file
diff --git a/_modules/cyto_dl/nn/vits/cross_mae.html b/_modules/cyto_dl/nn/vits/decoder.html
similarity index 73%
rename from _modules/cyto_dl/nn/vits/cross_mae.html
rename to _modules/cyto_dl/nn/vits/decoder.html
index 534dbb7f9..334210408 100644
--- a/_modules/cyto_dl/nn/vits/cross_mae.html
+++ b/_modules/cyto_dl/nn/vits/decoder.html
@@ -5,7 +5,7 @@
     <meta name="color-scheme" content="light dark"><link rel="index" title="Index" href="../../../../genindex.html" /><link rel="search" title="Search" href="../../../../search.html" />
 
     <!-- Generated with Sphinx 6.1.3 and Furo 2023.03.23 -->
-        <title>cyto_dl.nn.vits.cross_mae - cyto-dl 0.4.0 documentation</title>
+        <title>cyto_dl.nn.vits.decoder - cyto-dl 0.4.0 documentation</title>
       <link rel="stylesheet" type="text/css" href="../../../../_static/pygments.css" />
     <link rel="stylesheet" type="text/css" href="../../../../_static/styles/furo.css?digest=fad236701ea90a88636c2a8c73b44ae642ed2a53" />
     <link rel="stylesheet" type="text/css" href="../../../../_static/styles/furo-extensions.css?digest=30d1aed668e5c3a91c3e3bf6a60b675221979f0e" />
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
@@ -432,75 +439,78 @@
           </label>
         </div>
         <article role="main">
-          <h1>Source code for cyto_dl.nn.vits.cross_mae</h1><div class="highlight"><pre>
-<span></span><span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">List</span><span class="p">,</span> <span class="n">Optional</span>
+          <h1>Source code for cyto_dl.nn.vits.decoder</h1><div class="highlight"><pre>
+<span></span><span class="c1"># modified from https://github.com/IcarusWizard/MAE/blob/main/model.py#L124</span>
+
+<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">List</span><span class="p">,</span> <span class="n">Optional</span><span class="p">,</span> <span class="n">Union</span>
 
 <span class="kn">import</span> <span class="nn">torch</span>
 <span class="kn">import</span> <span class="nn">torch.nn</span> <span class="k">as</span> <span class="nn">nn</span>
-<span class="kn">from</span> <span class="nn">einops</span> <span class="kn">import</span> <span class="n">rearrange</span>
+<span class="kn">from</span> <span class="nn">einops</span> <span class="kn">import</span> <span class="n">rearrange</span><span class="p">,</span> <span class="n">repeat</span>
 <span class="kn">from</span> <span class="nn">einops.layers.torch</span> <span class="kn">import</span> <span class="n">Rearrange</span>
 <span class="kn">from</span> <span class="nn">timm.models.layers</span> <span class="kn">import</span> <span class="n">trunc_normal_</span>
+<span class="kn">from</span> <span class="nn">timm.models.vision_transformer</span> <span class="kn">import</span> <span class="n">Block</span>
 
 <span class="kn">from</span> <span class="nn">cyto_dl.nn.vits.blocks</span> <span class="kn">import</span> <span class="n">CrossAttentionBlock</span>
-<span class="kn">from</span> <span class="nn">cyto_dl.nn.vits.utils</span> <span class="kn">import</span> <span class="n">get_positional_embedding</span><span class="p">,</span> <span class="n">take_indexes</span>
+<span class="kn">from</span> <span class="nn">cyto_dl.nn.vits.utils</span> <span class="kn">import</span> <span class="p">(</span>
+    <span class="n">get_positional_embedding</span><span class="p">,</span>
+    <span class="n">match_tuple_dimensions</span><span class="p">,</span>
+    <span class="n">take_indexes</span><span class="p">,</span>
+<span class="p">)</span>
 
 
-<div class="viewcode-block" id="CrossMAE_Decoder"><a class="viewcode-back" href="../../../../cyto_dl.nn.vits.cross_mae.html#cyto_dl.nn.vits.cross_mae.CrossMAE_Decoder">[docs]</a><span class="k">class</span> <span class="nc">CrossMAE_Decoder</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
-<span class="w">    </span><span class="sd">&quot;&quot;&quot;Decoder inspired by [CrossMAE](https://crossmae.github.io/) where masked tokens only attend</span>
-<span class="sd">    to visible tokens.&quot;&quot;&quot;</span>
-
+<div class="viewcode-block" id="MAE_Decoder"><a class="viewcode-back" href="../../../../cyto_dl.nn.vits.decoder.html#cyto_dl.nn.vits.decoder.MAE_Decoder">[docs]</a><span class="k">class</span> <span class="nc">MAE_Decoder</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
     <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span>
         <span class="bp">self</span><span class="p">,</span>
-        <span class="n">num_patches</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span>
+        <span class="n">num_patches</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">]],</span>
         <span class="n">spatial_dims</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">3</span><span class="p">,</span>
-        <span class="n">base_patch_size</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">]]</span> <span class="o">=</span> <span class="p">[</span><span class="mi">4</span><span class="p">,</span> <span class="mi">8</span><span class="p">,</span> <span class="mi">8</span><span class="p">],</span>
+        <span class="n">patch_size</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">]]]</span> <span class="o">=</span> <span class="mi">4</span><span class="p">,</span>
         <span class="n">enc_dim</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">768</span><span class="p">,</span>
         <span class="n">emb_dim</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">192</span><span class="p">,</span>
         <span class="n">num_layer</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">4</span><span class="p">,</span>
         <span class="n">num_head</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">3</span><span class="p">,</span>
+        <span class="n">has_cls_token</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span>
         <span class="n">learnable_pos_embedding</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">True</span><span class="p">,</span>
     <span class="p">)</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
 <span class="w">        </span><span class="sd">&quot;&quot;&quot;</span>
 <span class="sd">        Parameters</span>
 <span class="sd">        ----------</span>
-<span class="sd">        num_patches: List[int]</span>
-<span class="sd">            Number of patches in each dimension</span>
-<span class="sd">        base_patch_size: Tuple[int]</span>
-<span class="sd">            Size of each patch</span>
+<span class="sd">        num_patches: List[int], int</span>
+<span class="sd">            Number of patches in each dimension. If int, the same number of patches is used for all dimensions.</span>
+<span class="sd">        patch_size: Tuple[int], int</span>
+<span class="sd">            Size of each patch. If int, the same patch size is used for all dimensions.</span>
 <span class="sd">        enc_dim: int</span>
 <span class="sd">            Dimension of encoder</span>
 <span class="sd">        emb_dim: int</span>
-<span class="sd">            Dimension of embedding</span>
+<span class="sd">            Dimension of decoder</span>
 <span class="sd">        num_layer: int</span>
 <span class="sd">            Number of transformer layers</span>
 <span class="sd">        num_head: int</span>
 <span class="sd">            Number of heads in transformer</span>
+<span class="sd">        has_cls_token: bool</span>
+<span class="sd">            Whether encoder features have a cls token</span>
 <span class="sd">        learnable_pos_embedding: bool</span>
-<span class="sd">            If True, learnable positional embeddings are used. If False, fixed sin/cos positional embeddings are used. Empirically, fixed positional embeddings work better for brightfield images.</span>
+<span class="sd">            If True, learnable positional embeddings are used. If False, fixed sin/cos positional embeddings. Empirically, fixed positional embeddings work better for brightfield images.</span>
 <span class="sd">        &quot;&quot;&quot;</span>
         <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>
+        <span class="n">num_patches</span><span class="p">,</span> <span class="n">patch_size</span> <span class="o">=</span> <span class="n">match_tuple_dimensions</span><span class="p">(</span><span class="n">spatial_dims</span><span class="p">,</span> <span class="p">[</span><span class="n">num_patches</span><span class="p">,</span> <span class="n">patch_size</span><span class="p">])</span>
 
-        <span class="bp">self</span><span class="o">.</span><span class="n">transformer</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">ParameterList</span><span class="p">(</span>
-            <span class="p">[</span>
-                <span class="n">CrossAttentionBlock</span><span class="p">(</span>
-                    <span class="n">encoder_dim</span><span class="o">=</span><span class="n">emb_dim</span><span class="p">,</span>
-                    <span class="n">decoder_dim</span><span class="o">=</span><span class="n">emb_dim</span><span class="p">,</span>
-                    <span class="n">num_heads</span><span class="o">=</span><span class="n">num_head</span><span class="p">,</span>
-                <span class="p">)</span>
-                <span class="k">for</span> <span class="n">_</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">num_layer</span><span class="p">)</span>
-            <span class="p">]</span>
-        <span class="p">)</span>
-        <span class="bp">self</span><span class="o">.</span><span class="n">decoder_norm</span> <span class="o">=</span> <span class="n">nn</span><span class="o">.</span><span class="n">LayerNorm</span><span class="p">(</span><span class="n">emb_dim</span><span class="p">)</span>
-        <span class="bp">self</span><span class="o">.</span><span class="n">projection_norm</span> <span class="o">=</span> <span class="n">nn</span><span class="o">.</span><span class="n">LayerNorm</span><span class="p">(</span><span class="n">emb_dim</span><span class="p">)</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">has_cls_token</span> <span class="o">=</span> <span class="n">has_cls_token</span>
 
+        <span class="bp">self</span><span class="o">.</span><span class="n">projection_norm</span> <span class="o">=</span> <span class="n">nn</span><span class="o">.</span><span class="n">LayerNorm</span><span class="p">(</span><span class="n">emb_dim</span><span class="p">)</span>
         <span class="bp">self</span><span class="o">.</span><span class="n">projection</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Linear</span><span class="p">(</span><span class="n">enc_dim</span><span class="p">,</span> <span class="n">emb_dim</span><span class="p">)</span>
         <span class="bp">self</span><span class="o">.</span><span class="n">mask_token</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Parameter</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="n">emb_dim</span><span class="p">))</span>
 
         <span class="bp">self</span><span class="o">.</span><span class="n">pos_embedding</span> <span class="o">=</span> <span class="n">get_positional_embedding</span><span class="p">(</span>
-            <span class="n">num_patches</span><span class="p">,</span> <span class="n">emb_dim</span><span class="p">,</span> <span class="n">learnable</span><span class="o">=</span><span class="n">learnable_pos_embedding</span>
+            <span class="n">num_patches</span><span class="p">,</span> <span class="n">emb_dim</span><span class="p">,</span> <span class="n">use_cls_token</span><span class="o">=</span><span class="n">has_cls_token</span><span class="p">,</span> <span class="n">learnable</span><span class="o">=</span><span class="n">learnable_pos_embedding</span>
         <span class="p">)</span>
 
-        <span class="bp">self</span><span class="o">.</span><span class="n">head</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Linear</span><span class="p">(</span><span class="n">emb_dim</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">prod</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">as_tensor</span><span class="p">(</span><span class="n">base_patch_size</span><span class="p">)))</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">transformer</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Sequential</span><span class="p">(</span>
+            <span class="o">*</span><span class="p">[</span><span class="n">Block</span><span class="p">(</span><span class="n">emb_dim</span><span class="p">,</span> <span class="n">num_head</span><span class="p">)</span> <span class="k">for</span> <span class="n">_</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">num_layer</span><span class="p">)]</span>
+        <span class="p">)</span>
+        <span class="n">out_dim</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">prod</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">as_tensor</span><span class="p">(</span><span class="n">patch_size</span><span class="p">))</span><span class="o">.</span><span class="n">item</span><span class="p">()</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">decoder_norm</span> <span class="o">=</span> <span class="n">nn</span><span class="o">.</span><span class="n">LayerNorm</span><span class="p">(</span><span class="n">emb_dim</span><span class="p">)</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">head</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Linear</span><span class="p">(</span><span class="n">emb_dim</span><span class="p">,</span> <span class="n">out_dim</span><span class="p">)</span>
         <span class="bp">self</span><span class="o">.</span><span class="n">num_patches</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">as_tensor</span><span class="p">(</span><span class="n">num_patches</span><span class="p">)</span>
 
         <span class="k">if</span> <span class="n">spatial_dims</span> <span class="o">==</span> <span class="mi">3</span><span class="p">:</span>
@@ -509,61 +519,144 @@ <h1>Source code for cyto_dl.nn.vits.cross_mae</h1><div class="highlight"><pre>
                 <span class="n">n_patch_z</span><span class="o">=</span><span class="n">num_patches</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span>
                 <span class="n">n_patch_y</span><span class="o">=</span><span class="n">num_patches</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span>
                 <span class="n">n_patch_x</span><span class="o">=</span><span class="n">num_patches</span><span class="p">[</span><span class="mi">2</span><span class="p">],</span>
-                <span class="n">patch_size_z</span><span class="o">=</span><span class="n">base_patch_size</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span>
-                <span class="n">patch_size_y</span><span class="o">=</span><span class="n">base_patch_size</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span>
-                <span class="n">patch_size_x</span><span class="o">=</span><span class="n">base_patch_size</span><span class="p">[</span><span class="mi">2</span><span class="p">],</span>
+                <span class="n">patch_size_z</span><span class="o">=</span><span class="n">patch_size</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span>
+                <span class="n">patch_size_y</span><span class="o">=</span><span class="n">patch_size</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span>
+                <span class="n">patch_size_x</span><span class="o">=</span><span class="n">patch_size</span><span class="p">[</span><span class="mi">2</span><span class="p">],</span>
             <span class="p">)</span>
         <span class="k">elif</span> <span class="n">spatial_dims</span> <span class="o">==</span> <span class="mi">2</span><span class="p">:</span>
             <span class="bp">self</span><span class="o">.</span><span class="n">patch2img</span> <span class="o">=</span> <span class="n">Rearrange</span><span class="p">(</span>
                 <span class="s2">&quot;(n_patch_y n_patch_x) b (c patch_size_y patch_size_x) -&gt;  b c (n_patch_y patch_size_y) (n_patch_x patch_size_x)&quot;</span><span class="p">,</span>
                 <span class="n">n_patch_y</span><span class="o">=</span><span class="n">num_patches</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span>
                 <span class="n">n_patch_x</span><span class="o">=</span><span class="n">num_patches</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span>
-                <span class="n">patch_size_y</span><span class="o">=</span><span class="n">base_patch_size</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span>
-                <span class="n">patch_size_x</span><span class="o">=</span><span class="n">base_patch_size</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span>
+                <span class="n">patch_size_y</span><span class="o">=</span><span class="n">patch_size</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span>
+                <span class="n">patch_size_x</span><span class="o">=</span><span class="n">patch_size</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span>
             <span class="p">)</span>
 
         <span class="bp">self</span><span class="o">.</span><span class="n">init_weight</span><span class="p">()</span>
 
-<div class="viewcode-block" id="CrossMAE_Decoder.init_weight"><a class="viewcode-back" href="../../../../cyto_dl.nn.vits.cross_mae.html#cyto_dl.nn.vits.cross_mae.CrossMAE_Decoder.init_weight">[docs]</a>    <span class="k">def</span> <span class="nf">init_weight</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
+<div class="viewcode-block" id="MAE_Decoder.init_weight"><a class="viewcode-back" href="../../../../cyto_dl.nn.vits.decoder.html#cyto_dl.nn.vits.decoder.MAE_Decoder.init_weight">[docs]</a>    <span class="k">def</span> <span class="nf">init_weight</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
         <span class="n">trunc_normal_</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">mask_token</span><span class="p">,</span> <span class="n">std</span><span class="o">=</span><span class="mf">0.02</span><span class="p">)</span></div>
 
-<div class="viewcode-block" id="CrossMAE_Decoder.forward"><a class="viewcode-back" href="../../../../cyto_dl.nn.vits.cross_mae.html#cyto_dl.nn.vits.cross_mae.CrossMAE_Decoder.forward">[docs]</a>    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">features</span><span class="p">,</span> <span class="n">forward_indexes</span><span class="p">,</span> <span class="n">backward_indexes</span><span class="p">):</span>
-        <span class="c1"># HACK TODO allow usage of multiple intermediate feature weights, this works when decoder is 0 layers</span>
+<div class="viewcode-block" id="MAE_Decoder.adjust_indices_for_cls"><a class="viewcode-back" href="../../../../cyto_dl.nn.vits.decoder.html#cyto_dl.nn.vits.decoder.MAE_Decoder.adjust_indices_for_cls">[docs]</a>    <span class="k">def</span> <span class="nf">adjust_indices_for_cls</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">indexes</span><span class="p">):</span>
+        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">has_cls_token</span><span class="p">:</span>
+            <span class="c1"># add all zeros to indices - this keeps the class token as the first index in the</span>
+            <span class="c1"># tensor. We also have to add 1 to all the indices to account for the zeros we added</span>
+            <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">cat</span><span class="p">(</span>
+                <span class="p">[</span>
+                    <span class="n">torch</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="n">indexes</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class="n">device</span><span class="o">=</span><span class="n">indexes</span><span class="o">.</span><span class="n">device</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">torch</span><span class="o">.</span><span class="n">long</span><span class="p">),</span>
+                    <span class="n">indexes</span> <span class="o">+</span> <span class="mi">1</span><span class="p">,</span>
+                <span class="p">],</span>
+                <span class="n">dim</span><span class="o">=</span><span class="mi">0</span><span class="p">,</span>
+            <span class="p">)</span>
+        <span class="k">return</span> <span class="n">indexes</span></div>
+
+<div class="viewcode-block" id="MAE_Decoder.add_mask_tokens"><a class="viewcode-back" href="../../../../cyto_dl.nn.vits.decoder.html#cyto_dl.nn.vits.decoder.MAE_Decoder.add_mask_tokens">[docs]</a>    <span class="k">def</span> <span class="nf">add_mask_tokens</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">features</span><span class="p">,</span> <span class="n">backward_indexes</span><span class="p">):</span>
+        <span class="c1"># fill in deleted masked regions with mask token</span>
+        <span class="n">num_visible_tokens</span><span class="p">,</span> <span class="n">B</span><span class="p">,</span> <span class="n">_</span> <span class="o">=</span> <span class="n">features</span><span class="o">.</span><span class="n">shape</span>
+        <span class="c1"># total number of tokens - number of visible tokens</span>
+        <span class="n">num_mask_tokens</span> <span class="o">=</span> <span class="n">backward_indexes</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">-</span> <span class="n">num_visible_tokens</span>
+        <span class="n">mask_tokens</span> <span class="o">=</span> <span class="n">repeat</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">mask_token</span><span class="p">,</span> <span class="s2">&quot;1 1 c -&gt; t b c&quot;</span><span class="p">,</span> <span class="n">t</span><span class="o">=</span><span class="n">num_mask_tokens</span><span class="p">,</span> <span class="n">b</span><span class="o">=</span><span class="n">B</span><span class="p">)</span>
+        <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">cat</span><span class="p">([</span><span class="n">features</span><span class="p">,</span> <span class="n">mask_tokens</span><span class="p">],</span> <span class="n">dim</span><span class="o">=</span><span class="mi">0</span><span class="p">)</span></div>
+
+<div class="viewcode-block" id="MAE_Decoder.forward"><a class="viewcode-back" href="../../../../cyto_dl.nn.vits.decoder.html#cyto_dl.nn.vits.decoder.MAE_Decoder.forward">[docs]</a>    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">features</span><span class="p">,</span> <span class="n">forward_indexes</span><span class="p">,</span> <span class="n">backward_indexes</span><span class="p">):</span>
+        <span class="c1"># project from encoder dimension to decoder dimension</span>
+        <span class="n">features</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">projection_norm</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">projection</span><span class="p">(</span><span class="n">features</span><span class="p">))</span>
+
+        <span class="n">backward_indexes</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">adjust_indices_for_cls</span><span class="p">(</span><span class="n">backward_indexes</span><span class="p">)</span>
+
+        <span class="n">features</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">add_mask_tokens</span><span class="p">(</span><span class="n">features</span><span class="p">,</span> <span class="n">backward_indexes</span><span class="p">)</span>
+
+        <span class="c1"># unshuffle to original positions</span>
+        <span class="n">features</span> <span class="o">=</span> <span class="n">take_indexes</span><span class="p">(</span><span class="n">features</span><span class="p">,</span> <span class="n">backward_indexes</span><span class="p">)</span>
+        <span class="n">features</span> <span class="o">=</span> <span class="n">features</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">pos_embedding</span>
+
+        <span class="c1"># decode</span>
+        <span class="n">features</span> <span class="o">=</span> <span class="n">rearrange</span><span class="p">(</span><span class="n">features</span><span class="p">,</span> <span class="s2">&quot;t b c -&gt; b t c&quot;</span><span class="p">)</span>
+        <span class="n">features</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">transformer</span><span class="p">(</span><span class="n">features</span><span class="p">)</span>
+        <span class="n">features</span> <span class="o">=</span> <span class="n">rearrange</span><span class="p">(</span><span class="n">features</span><span class="p">,</span> <span class="s2">&quot;b t c -&gt; t b c&quot;</span><span class="p">)</span>
+
+        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">has_cls_token</span><span class="p">:</span>
+            <span class="n">features</span> <span class="o">=</span> <span class="n">features</span><span class="p">[</span><span class="mi">1</span><span class="p">:]</span>  <span class="c1"># remove global feature</span>
+
+        <span class="c1"># (npatches x npatches x npatches) b (emb dim) -&gt; (npatches* npatches * npatches) b (z y x)</span>
+        <span class="n">patches</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">head</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">decoder_norm</span><span class="p">(</span><span class="n">features</span><span class="p">))</span>
+
+        <span class="c1"># patches to image</span>
+        <span class="n">img</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">patch2img</span><span class="p">(</span><span class="n">patches</span><span class="p">)</span>
+        <span class="k">return</span> <span class="n">img</span></div></div>
+
+
+<div class="viewcode-block" id="CrossMAE_Decoder"><a class="viewcode-back" href="../../../../cyto_dl.nn.vits.decoder.html#cyto_dl.nn.vits.decoder.CrossMAE_Decoder">[docs]</a><span class="k">class</span> <span class="nc">CrossMAE_Decoder</span><span class="p">(</span><span class="n">MAE_Decoder</span><span class="p">):</span>
+<span class="w">    </span><span class="sd">&quot;&quot;&quot;Decoder inspired by [CrossMAE](https://crossmae.github.io/) where masked tokens only attend</span>
+<span class="sd">    to visible tokens.&quot;&quot;&quot;</span>
+
+    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span>
+        <span class="bp">self</span><span class="p">,</span>
+        <span class="n">num_patches</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">]],</span>
+        <span class="n">spatial_dims</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">3</span><span class="p">,</span>
+        <span class="n">patch_size</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">]]]</span> <span class="o">=</span> <span class="mi">4</span><span class="p">,</span>
+        <span class="n">enc_dim</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">768</span><span class="p">,</span>
+        <span class="n">emb_dim</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">192</span><span class="p">,</span>
+        <span class="n">num_layer</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">4</span><span class="p">,</span>
+        <span class="n">num_head</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">3</span><span class="p">,</span>
+        <span class="n">has_cls_token</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">True</span><span class="p">,</span>
+        <span class="n">learnable_pos_embedding</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">True</span><span class="p">,</span>
+    <span class="p">)</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
+<span class="w">        </span><span class="sd">&quot;&quot;&quot;</span>
+<span class="sd">        Parameters</span>
+<span class="sd">        ----------</span>
+<span class="sd">        num_patches: List[int], int</span>
+<span class="sd">            Number of patches in each dimension. If int, the same number of patches is used for all dimensions.</span>
+<span class="sd">        patch_size: Tuple[int]</span>
+<span class="sd">            Size of each patch in each dimension. If int, the same patch size is used for all dimensions.</span>
+<span class="sd">        enc_dim: int</span>
+<span class="sd">            Dimension of encoder</span>
+<span class="sd">        emb_dim: int</span>
+<span class="sd">            Dimension of embedding</span>
+<span class="sd">        num_layer: int</span>
+<span class="sd">            Number of transformer layers</span>
+<span class="sd">        num_head: int</span>
+<span class="sd">            Number of heads in transformer</span>
+<span class="sd">        has_cls_token: bool</span>
+<span class="sd">            Whether encoder features have a cls token</span>
+<span class="sd">        learnable_pos_embedding: bool</span>
+<span class="sd">            If True, learnable positional embeddings are used. If False, fixed sin/cos positional embeddings are used. Empirically, fixed positional embeddings work better for brightfield images.</span>
+<span class="sd">        &quot;&quot;&quot;</span>
+        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span>
+            <span class="n">num_patches</span><span class="p">,</span>
+            <span class="n">spatial_dims</span><span class="p">,</span>
+            <span class="n">patch_size</span><span class="p">,</span>
+            <span class="n">enc_dim</span><span class="p">,</span>
+            <span class="n">emb_dim</span><span class="p">,</span>
+            <span class="n">num_layer</span><span class="p">,</span>
+            <span class="n">num_head</span><span class="p">,</span>
+            <span class="n">has_cls_token</span><span class="p">,</span>
+            <span class="n">learnable_pos_embedding</span><span class="p">,</span>
+        <span class="p">)</span>
+
+        <span class="bp">self</span><span class="o">.</span><span class="n">transformer</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">ParameterList</span><span class="p">(</span>
+            <span class="p">[</span>
+                <span class="n">CrossAttentionBlock</span><span class="p">(</span>
+                    <span class="n">encoder_dim</span><span class="o">=</span><span class="n">emb_dim</span><span class="p">,</span>
+                    <span class="n">decoder_dim</span><span class="o">=</span><span class="n">emb_dim</span><span class="p">,</span>
+                    <span class="n">num_heads</span><span class="o">=</span><span class="n">num_head</span><span class="p">,</span>
+                <span class="p">)</span>
+                <span class="k">for</span> <span class="n">_</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">num_layer</span><span class="p">)</span>
+            <span class="p">]</span>
+        <span class="p">)</span>
+
+<div class="viewcode-block" id="CrossMAE_Decoder.forward"><a class="viewcode-back" href="../../../../cyto_dl.nn.vits.decoder.html#cyto_dl.nn.vits.decoder.CrossMAE_Decoder.forward">[docs]</a>    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">features</span><span class="p">,</span> <span class="n">forward_indexes</span><span class="p">,</span> <span class="n">backward_indexes</span><span class="p">):</span>
+        <span class="c1"># HACK TODO allow usage of multiple intermediate feature weights, this works when decoder is 1 layer</span>
         <span class="c1"># features can be n t b c (if intermediate feature weighter used) or t b c if not</span>
         <span class="n">features</span> <span class="o">=</span> <span class="n">features</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">features</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span> <span class="o">==</span> <span class="mi">4</span> <span class="k">else</span> <span class="n">features</span>
         <span class="n">T</span><span class="p">,</span> <span class="n">B</span><span class="p">,</span> <span class="n">C</span> <span class="o">=</span> <span class="n">features</span><span class="o">.</span><span class="n">shape</span>
         <span class="c1"># we could do cross attention between decoder_dim queries and encoder_dim features, but it seems to work fine having both at decoder_dim for now</span>
         <span class="n">features</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">projection_norm</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">projection</span><span class="p">(</span><span class="n">features</span><span class="p">))</span>
 
-        <span class="c1"># add cls token</span>
-        <span class="n">backward_indexes</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">cat</span><span class="p">(</span>
-            <span class="p">[</span>
-                <span class="n">torch</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span>
-                    <span class="mi">1</span><span class="p">,</span> <span class="n">backward_indexes</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class="n">device</span><span class="o">=</span><span class="n">backward_indexes</span><span class="o">.</span><span class="n">device</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">torch</span><span class="o">.</span><span class="n">long</span>
-                <span class="p">),</span>
-                <span class="n">backward_indexes</span> <span class="o">+</span> <span class="mi">1</span><span class="p">,</span>
-            <span class="p">],</span>
-            <span class="n">dim</span><span class="o">=</span><span class="mi">0</span><span class="p">,</span>
-        <span class="p">)</span>
-        <span class="n">forward_indexes</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">cat</span><span class="p">(</span>
-            <span class="p">[</span>
-                <span class="n">torch</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span>
-                    <span class="mi">1</span><span class="p">,</span> <span class="n">forward_indexes</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class="n">device</span><span class="o">=</span><span class="n">forward_indexes</span><span class="o">.</span><span class="n">device</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">torch</span><span class="o">.</span><span class="n">long</span>
-                <span class="p">),</span>
-                <span class="n">forward_indexes</span> <span class="o">+</span> <span class="mi">1</span><span class="p">,</span>
-            <span class="p">],</span>
-            <span class="n">dim</span><span class="o">=</span><span class="mi">0</span><span class="p">,</span>
-        <span class="p">)</span>
-        <span class="c1"># fill in masked regions</span>
-        <span class="n">features</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">cat</span><span class="p">(</span>
-            <span class="p">[</span>
-                <span class="n">features</span><span class="p">,</span>
-                <span class="bp">self</span><span class="o">.</span><span class="n">mask_token</span><span class="o">.</span><span class="n">expand</span><span class="p">(</span>
-                    <span class="n">backward_indexes</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">-</span> <span class="n">features</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">features</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class="o">-</span><span class="mi">1</span>
-                <span class="p">),</span>
-            <span class="p">],</span>
-            <span class="n">dim</span><span class="o">=</span><span class="mi">0</span><span class="p">,</span>
-        <span class="p">)</span>
+        <span class="n">backward_indexes</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">adjust_indices_for_cls</span><span class="p">(</span><span class="n">backward_indexes</span><span class="p">)</span>
+        <span class="n">forward_indexes</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">adjust_indices_for_cls</span><span class="p">(</span><span class="n">forward_indexes</span><span class="p">)</span>
+
+        <span class="n">features</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">add_mask_tokens</span><span class="p">(</span><span class="n">features</span><span class="p">,</span> <span class="n">backward_indexes</span><span class="p">)</span>
 
         <span class="c1"># unshuffle to original positions for positional embedding so we can do cross attention during decoding</span>
         <span class="n">features</span> <span class="o">=</span> <span class="n">take_indexes</span><span class="p">(</span><span class="n">features</span><span class="p">,</span> <span class="n">backward_indexes</span><span class="p">)</span>
@@ -590,7 +683,8 @@ <h1>Source code for cyto_dl.nn.vits.cross_mae</h1><div class="highlight"><pre>
         <span class="n">patches</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">cat</span><span class="p">(</span>
             <span class="p">[</span>
                 <span class="n">torch</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span>
-                    <span class="p">(</span><span class="n">T</span> <span class="o">-</span> <span class="mi">1</span><span class="p">,</span> <span class="n">B</span><span class="p">,</span> <span class="n">patches</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">]),</span>
+                    <span class="c1"># T-1 accounts for cls token</span>
+                    <span class="p">(</span><span class="n">T</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">has_cls_token</span><span class="p">,</span> <span class="n">B</span><span class="p">,</span> <span class="n">patches</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">]),</span>
                     <span class="n">requires_grad</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
                     <span class="n">device</span><span class="o">=</span><span class="n">patches</span><span class="o">.</span><span class="n">device</span><span class="p">,</span>
                     <span class="n">dtype</span><span class="o">=</span><span class="n">patches</span><span class="o">.</span><span class="n">dtype</span><span class="p">,</span>
@@ -599,7 +693,11 @@ <h1>Source code for cyto_dl.nn.vits.cross_mae</h1><div class="highlight"><pre>
             <span class="p">],</span>
             <span class="n">dim</span><span class="o">=</span><span class="mi">0</span><span class="p">,</span>
         <span class="p">)</span>
-        <span class="n">patches</span> <span class="o">=</span> <span class="n">take_indexes</span><span class="p">(</span><span class="n">patches</span><span class="p">,</span> <span class="n">backward_indexes</span><span class="p">[</span><span class="mi">1</span><span class="p">:]</span> <span class="o">-</span> <span class="mi">1</span><span class="p">)</span>
+        <span class="n">patches</span> <span class="o">=</span> <span class="p">(</span>
+            <span class="n">take_indexes</span><span class="p">(</span><span class="n">patches</span><span class="p">,</span> <span class="n">backward_indexes</span><span class="p">[</span><span class="mi">1</span><span class="p">:]</span> <span class="o">-</span> <span class="mi">1</span><span class="p">)</span>
+            <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">has_cls_token</span>
+            <span class="k">else</span> <span class="n">take_indexes</span><span class="p">(</span><span class="n">patches</span><span class="p">,</span> <span class="n">backward_indexes</span><span class="p">)</span>
+        <span class="p">)</span>
         <span class="c1"># patches to image</span>
         <span class="n">img</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">patch2img</span><span class="p">(</span><span class="n">patches</span><span class="p">)</span>
         <span class="k">return</span> <span class="n">img</span></div></div>
diff --git a/_modules/cyto_dl/nn/vits/encoder.html b/_modules/cyto_dl/nn/vits/encoder.html
new file mode 100644
index 000000000..a1c17ffd8
--- /dev/null
+++ b/_modules/cyto_dl/nn/vits/encoder.html
@@ -0,0 +1,863 @@
+<!doctype html>
+<html class="no-js" lang="en">
+  <head><meta charset="utf-8"/>
+    <meta name="viewport" content="width=device-width,initial-scale=1"/>
+    <meta name="color-scheme" content="light dark"><link rel="index" title="Index" href="../../../../genindex.html" /><link rel="search" title="Search" href="../../../../search.html" />
+
+    <!-- Generated with Sphinx 6.1.3 and Furo 2023.03.23 -->
+        <title>cyto_dl.nn.vits.encoder - cyto-dl 0.4.0 documentation</title>
+      <link rel="stylesheet" type="text/css" href="../../../../_static/pygments.css" />
+    <link rel="stylesheet" type="text/css" href="../../../../_static/styles/furo.css?digest=fad236701ea90a88636c2a8c73b44ae642ed2a53" />
+    <link rel="stylesheet" type="text/css" href="../../../../_static/styles/furo-extensions.css?digest=30d1aed668e5c3a91c3e3bf6a60b675221979f0e" />
+    
+    
+
+
+<style>
+  body {
+    --color-code-background: #eeffcc;
+  --color-code-foreground: black;
+  
+  }
+  @media not print {
+    body[data-theme="dark"] {
+      --color-code-background: #202020;
+  --color-code-foreground: #d0d0d0;
+  
+    }
+    @media (prefers-color-scheme: dark) {
+      body:not([data-theme="light"]) {
+        --color-code-background: #202020;
+  --color-code-foreground: #d0d0d0;
+  
+      }
+    }
+  }
+</style></head>
+  <body>
+    
+    <script>
+      document.body.dataset.theme = localStorage.getItem("theme") || "auto";
+    </script>
+    
+
+<svg xmlns="http://www.w3.org/2000/svg" style="display: none;">
+  <symbol id="svg-toc" viewBox="0 0 24 24">
+    <title>Contents</title>
+    <svg stroke="currentColor" fill="currentColor" stroke-width="0" viewBox="0 0 1024 1024">
+      <path d="M408 442h480c4.4 0 8-3.6 8-8v-56c0-4.4-3.6-8-8-8H408c-4.4 0-8 3.6-8 8v56c0 4.4 3.6 8 8 8zm-8 204c0 4.4 3.6 8 8 8h480c4.4 0 8-3.6 8-8v-56c0-4.4-3.6-8-8-8H408c-4.4 0-8 3.6-8 8v56zm504-486H120c-4.4 0-8 3.6-8 8v56c0 4.4 3.6 8 8 8h784c4.4 0 8-3.6 8-8v-56c0-4.4-3.6-8-8-8zm0 632H120c-4.4 0-8 3.6-8 8v56c0 4.4 3.6 8 8 8h784c4.4 0 8-3.6 8-8v-56c0-4.4-3.6-8-8-8zM115.4 518.9L271.7 642c5.8 4.6 14.4.5 14.4-6.9V388.9c0-7.4-8.5-11.5-14.4-6.9L115.4 505.1a8.74 8.74 0 0 0 0 13.8z"/>
+    </svg>
+  </symbol>
+  <symbol id="svg-menu" viewBox="0 0 24 24">
+    <title>Menu</title>
+    <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
+      stroke-width="2" stroke-linecap="round" stroke-linejoin="round" class="feather-menu">
+      <line x1="3" y1="12" x2="21" y2="12"></line>
+      <line x1="3" y1="6" x2="21" y2="6"></line>
+      <line x1="3" y1="18" x2="21" y2="18"></line>
+    </svg>
+  </symbol>
+  <symbol id="svg-arrow-right" viewBox="0 0 24 24">
+    <title>Expand</title>
+    <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
+      stroke-width="2" stroke-linecap="round" stroke-linejoin="round" class="feather-chevron-right">
+      <polyline points="9 18 15 12 9 6"></polyline>
+    </svg>
+  </symbol>
+  <symbol id="svg-sun" viewBox="0 0 24 24">
+    <title>Light mode</title>
+    <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
+      stroke-width="1.5" stroke-linecap="round" stroke-linejoin="round" class="feather-sun">
+      <circle cx="12" cy="12" r="5"></circle>
+      <line x1="12" y1="1" x2="12" y2="3"></line>
+      <line x1="12" y1="21" x2="12" y2="23"></line>
+      <line x1="4.22" y1="4.22" x2="5.64" y2="5.64"></line>
+      <line x1="18.36" y1="18.36" x2="19.78" y2="19.78"></line>
+      <line x1="1" y1="12" x2="3" y2="12"></line>
+      <line x1="21" y1="12" x2="23" y2="12"></line>
+      <line x1="4.22" y1="19.78" x2="5.64" y2="18.36"></line>
+      <line x1="18.36" y1="5.64" x2="19.78" y2="4.22"></line>
+    </svg>
+  </symbol>
+  <symbol id="svg-moon" viewBox="0 0 24 24">
+    <title>Dark mode</title>
+    <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
+      stroke-width="1.5" stroke-linecap="round" stroke-linejoin="round" class="icon-tabler-moon">
+      <path stroke="none" d="M0 0h24v24H0z" fill="none" />
+      <path d="M12 3c.132 0 .263 0 .393 0a7.5 7.5 0 0 0 7.92 12.446a9 9 0 1 1 -8.313 -12.454z" />
+    </svg>
+  </symbol>
+  <symbol id="svg-sun-half" viewBox="0 0 24 24">
+    <title>Auto light/dark mode</title>
+    <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
+      stroke-width="1.5" stroke-linecap="round" stroke-linejoin="round" class="icon-tabler-shadow">
+      <path stroke="none" d="M0 0h24v24H0z" fill="none"/>
+      <circle cx="12" cy="12" r="9" />
+      <path d="M13 12h5" />
+      <path d="M13 15h4" />
+      <path d="M13 18h1" />
+      <path d="M13 9h4" />
+      <path d="M13 6h1" />
+    </svg>
+  </symbol>
+</svg>
+
+<input type="checkbox" class="sidebar-toggle" name="__navigation" id="__navigation">
+<input type="checkbox" class="sidebar-toggle" name="__toc" id="__toc">
+<label class="overlay sidebar-overlay" for="__navigation">
+  <div class="visually-hidden">Hide navigation sidebar</div>
+</label>
+<label class="overlay toc-overlay" for="__toc">
+  <div class="visually-hidden">Hide table of contents sidebar</div>
+</label>
+
+
+
+<div class="page">
+  <header class="mobile-header">
+    <div class="header-left">
+      <label class="nav-overlay-icon" for="__navigation">
+        <div class="visually-hidden">Toggle site navigation sidebar</div>
+        <i class="icon"><svg><use href="#svg-menu"></use></svg></i>
+      </label>
+    </div>
+    <div class="header-center">
+      <a href="../../../../index.html"><div class="brand">cyto-dl 0.4.0 documentation</div></a>
+    </div>
+    <div class="header-right">
+      <div class="theme-toggle-container theme-toggle-header">
+        <button class="theme-toggle">
+          <div class="visually-hidden">Toggle Light / Dark / Auto color theme</div>
+          <svg class="theme-icon-when-auto"><use href="#svg-sun-half"></use></svg>
+          <svg class="theme-icon-when-dark"><use href="#svg-moon"></use></svg>
+          <svg class="theme-icon-when-light"><use href="#svg-sun"></use></svg>
+        </button>
+      </div>
+      <label class="toc-overlay-icon toc-header-icon no-toc" for="__toc">
+        <div class="visually-hidden">Toggle table of contents sidebar</div>
+        <i class="icon"><svg><use href="#svg-toc"></use></svg></i>
+      </label>
+    </div>
+  </header>
+  <aside class="sidebar-drawer">
+    <div class="sidebar-container">
+      
+      <div class="sidebar-sticky"><a class="sidebar-brand" href="../../../../index.html">
+  
+  
+  <span class="sidebar-brand-text">cyto-dl 0.4.0 documentation</span>
+  
+</a><form class="sidebar-search-container" method="get" action="../../../../search.html" role="search">
+  <input class="sidebar-search" placeholder="Search" name="q" aria-label="Search">
+  <input type="hidden" name="check_keywords" value="yes">
+  <input type="hidden" name="area" value="default">
+</form>
+<div id="searchbox"></div><div class="sidebar-scroll"><div class="sidebar-tree">
+  <p class="caption" role="heading"><span class="caption-text">Contents:</span></p>
+<ul>
+<li class="toctree-l1"><a class="reference internal" href="../../../../index.html">Home</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../../../installation.html">Installation</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../../../using_examples.html">Modifying the Example Configs</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../../../api.html">Using the API</a></li>
+<li class="toctree-l1 has-children"><a class="reference internal" href="../../../../modules.html">Package modules</a><input class="toctree-checkbox" id="toctree-checkbox-1" name="toctree-checkbox-1" role="switch" type="checkbox"/><label for="toctree-checkbox-1"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l2 has-children"><a class="reference internal" href="../../../../cyto_dl.html">cyto_dl package</a><input class="toctree-checkbox" id="toctree-checkbox-2" name="toctree-checkbox-2" role="switch" type="checkbox"/><label for="toctree-checkbox-2"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.api.html">cyto_dl.api package</a><input class="toctree-checkbox" id="toctree-checkbox-3" name="toctree-checkbox-3" role="switch" type="checkbox"/><label for="toctree-checkbox-3"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.api.cyto_dl_model.html">cyto_dl.api.cyto_dl_model package</a><input class="toctree-checkbox" id="toctree-checkbox-4" name="toctree-checkbox-4" role="switch" type="checkbox"/><label for="toctree-checkbox-4"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.api.cyto_dl_model.cyto_dl_base_model.html">cyto_dl.api.cyto_dl_model.cyto_dl_base_model module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.api.cyto_dl_model.segmentation_plugin_model.html">cyto_dl.api.cyto_dl_model.segmentation_plugin_model module</a></li>
+</ul>
+</li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.api.data.html">cyto_dl.api.data module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.api.model.html">cyto_dl.api.model module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.callbacks.html">cyto_dl.callbacks package</a><input class="toctree-checkbox" id="toctree-checkbox-5" name="toctree-checkbox-5" role="switch" type="checkbox"/><label for="toctree-checkbox-5"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.callbacks.callback_utils.html">cyto_dl.callbacks.callback_utils module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.callbacks.csv_saver.html">cyto_dl.callbacks.csv_saver module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.callbacks.image_saver.html">cyto_dl.callbacks.image_saver module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.callbacks.latent_walk.html">cyto_dl.callbacks.latent_walk module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.callbacks.layer_freeze.html">cyto_dl.callbacks.layer_freeze module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.callbacks.model_utils.html">cyto_dl.callbacks.model_utils module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.callbacks.outlier_detection.html">cyto_dl.callbacks.outlier_detection module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.dataframe.html">cyto_dl.dataframe package</a><input class="toctree-checkbox" id="toctree-checkbox-6" name="toctree-checkbox-6" role="switch" type="checkbox"/><label for="toctree-checkbox-6"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.dataframe.transforms.html">cyto_dl.dataframe.transforms package</a><input class="toctree-checkbox" id="toctree-checkbox-7" name="toctree-checkbox-7" role="switch" type="checkbox"/><label for="toctree-checkbox-7"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.dataframe.transforms.filter.html">cyto_dl.dataframe.transforms.filter module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.dataframe.transforms.group_cols.html">cyto_dl.dataframe.transforms.group_cols module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.dataframe.transforms.misc.html">cyto_dl.dataframe.transforms.misc module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.dataframe.transforms.split.html">cyto_dl.dataframe.transforms.split module</a></li>
+</ul>
+</li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.dataframe.readers.html">cyto_dl.dataframe.readers module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.datamodules.html">cyto_dl.datamodules package</a><input class="toctree-checkbox" id="toctree-checkbox-8" name="toctree-checkbox-8" role="switch" type="checkbox"/><label for="toctree-checkbox-8"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.datamodules.dataframe.html">cyto_dl.datamodules.dataframe package</a><input class="toctree-checkbox" id="toctree-checkbox-9" name="toctree-checkbox-9" role="switch" type="checkbox"/><label for="toctree-checkbox-9"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.datamodules.dataframe.dataframe_datamodule.html">cyto_dl.datamodules.dataframe.dataframe_datamodule module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.datamodules.dataframe.grouped_dataframe_datamodule.html">cyto_dl.datamodules.dataframe.grouped_dataframe_datamodule module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.datamodules.dataframe.utils.html">cyto_dl.datamodules.dataframe.utils module</a></li>
+</ul>
+</li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.datamodules.array.html">cyto_dl.datamodules.array module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.datamodules.data_dict.html">cyto_dl.datamodules.data_dict module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.datamodules.folder.html">cyto_dl.datamodules.folder module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.datamodules.multidim_image.html">cyto_dl.datamodules.multidim_image module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.datamodules.smartcache.html">cyto_dl.datamodules.smartcache module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.datamodules.torchvision.html">cyto_dl.datamodules.torchvision module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.image.html">cyto_dl.image package</a><input class="toctree-checkbox" id="toctree-checkbox-10" name="toctree-checkbox-10" role="switch" type="checkbox"/><label for="toctree-checkbox-10"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.image.io.html">cyto_dl.image.io package</a><input class="toctree-checkbox" id="toctree-checkbox-11" name="toctree-checkbox-11" role="switch" type="checkbox"/><label for="toctree-checkbox-11"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.image.io.aicsimage_loader.html">cyto_dl.image.io.aicsimage_loader module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.image.io.monai_bio_reader.html">cyto_dl.image.io.monai_bio_reader module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.image.io.numpy_reader.html">cyto_dl.image.io.numpy_reader module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.image.io.ome_zarr_reader.html">cyto_dl.image.io.ome_zarr_reader module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.image.io.polygon_loader.html">cyto_dl.image.io.polygon_loader module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.image.io.skimage_reader.html">cyto_dl.image.io.skimage_reader module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.image.transforms.html">cyto_dl.image.transforms package</a><input class="toctree-checkbox" id="toctree-checkbox-12" name="toctree-checkbox-12" role="switch" type="checkbox"/><label for="toctree-checkbox-12"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.image.transforms.add_meta.html">cyto_dl.image.transforms.add_meta module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.image.transforms.bright_sampler.html">cyto_dl.image.transforms.bright_sampler module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.image.transforms.clip.html">cyto_dl.image.transforms.clip module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.image.transforms.contrastadjust.html">cyto_dl.image.transforms.contrastadjust module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.image.transforms.generate_jepa_masks.html">cyto_dl.image.transforms.generate_jepa_masks module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.image.transforms.merge.html">cyto_dl.image.transforms.merge module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.image.transforms.multiscale_cropper.html">cyto_dl.image.transforms.multiscale_cropper module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.image.transforms.pad.html">cyto_dl.image.transforms.pad module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.image.transforms.project.html">cyto_dl.image.transforms.project module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.image.transforms.rotation_mask_transform.html">cyto_dl.image.transforms.rotation_mask_transform module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.image.transforms.save.html">cyto_dl.image.transforms.save module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.image.transforms.so2_random_rotation.html">cyto_dl.image.transforms.so2_random_rotation module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.image.transforms.track_transforms.html">cyto_dl.image.transforms.track_transforms module</a></li>
+</ul>
+</li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.loggers.html">cyto_dl.loggers package</a><input class="toctree-checkbox" id="toctree-checkbox-13" name="toctree-checkbox-13" role="switch" type="checkbox"/><label for="toctree-checkbox-13"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.loggers.mlflow.html">cyto_dl.loggers.mlflow module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.models.html">cyto_dl.models package</a><input class="toctree-checkbox" id="toctree-checkbox-14" name="toctree-checkbox-14" role="switch" type="checkbox"/><label for="toctree-checkbox-14"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.models.classification.html">cyto_dl.models.classification package</a><input class="toctree-checkbox" id="toctree-checkbox-15" name="toctree-checkbox-15" role="switch" type="checkbox"/><label for="toctree-checkbox-15"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.models.classification.classification.html">cyto_dl.models.classification.classification module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.models.classification.timepoint_classification.html">cyto_dl.models.classification.timepoint_classification module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.models.contrastive.html">cyto_dl.models.contrastive package</a><input class="toctree-checkbox" id="toctree-checkbox-16" name="toctree-checkbox-16" role="switch" type="checkbox"/><label for="toctree-checkbox-16"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.models.contrastive.contrastive.html">cyto_dl.models.contrastive.contrastive module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.models.handlers.html">cyto_dl.models.handlers package</a><input class="toctree-checkbox" id="toctree-checkbox-17" name="toctree-checkbox-17" role="switch" type="checkbox"/><label for="toctree-checkbox-17"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.models.handlers.base_handler.html">cyto_dl.models.handlers.base_handler module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.models.handlers.image_handler.html">cyto_dl.models.handlers.image_handler module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.models.handlers.load_image_patch.html">cyto_dl.models.handlers.load_image_patch module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.models.im2im.html">cyto_dl.models.im2im package</a><input class="toctree-checkbox" id="toctree-checkbox-18" name="toctree-checkbox-18" role="switch" type="checkbox"/><label for="toctree-checkbox-18"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.models.im2im.utils.html">cyto_dl.models.im2im.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-19" name="toctree-checkbox-19" role="switch" type="checkbox"/><label for="toctree-checkbox-19"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.models.im2im.utils.postprocessing.html">cyto_dl.models.im2im.utils.postprocessing package</a><input class="toctree-checkbox" id="toctree-checkbox-20" name="toctree-checkbox-20" role="switch" type="checkbox"/><label for="toctree-checkbox-20"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.models.im2im.utils.postprocessing.act_thresh_label.html">cyto_dl.models.im2im.utils.postprocessing.act_thresh_label module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.models.im2im.utils.postprocessing.auto_thresh.html">cyto_dl.models.im2im.utils.postprocessing.auto_thresh module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.models.im2im.utils.postprocessing.dict_to_im.html">cyto_dl.models.im2im.utils.postprocessing.dict_to_im module</a></li>
+</ul>
+</li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.models.im2im.utils.instance_seg.html">cyto_dl.models.im2im.utils.instance_seg module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.models.im2im.utils.noise_annealer.html">cyto_dl.models.im2im.utils.noise_annealer module</a></li>
+</ul>
+</li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.models.im2im.gan.html">cyto_dl.models.im2im.gan module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.models.im2im.multi_task.html">cyto_dl.models.im2im.multi_task module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.models.jepa.html">cyto_dl.models.jepa package</a><input class="toctree-checkbox" id="toctree-checkbox-21" name="toctree-checkbox-21" role="switch" type="checkbox"/><label for="toctree-checkbox-21"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.models.jepa.ijepa.html">cyto_dl.models.jepa.ijepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.models.jepa.iwm.html">cyto_dl.models.jepa.iwm module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.models.jepa.jepa_base.html">cyto_dl.models.jepa.jepa_base module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.models.utils.html">cyto_dl.models.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-22" name="toctree-checkbox-22" role="switch" type="checkbox"/><label for="toctree-checkbox-22"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.models.utils.mlflow.html">cyto_dl.models.utils.mlflow module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.models.utils.utils.html">cyto_dl.models.utils.utils module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.models.vae.html">cyto_dl.models.vae package</a><input class="toctree-checkbox" id="toctree-checkbox-23" name="toctree-checkbox-23" role="switch" type="checkbox"/><label for="toctree-checkbox-23"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.models.vae.o2_spharm_vae.html">cyto_dl.models.vae.o2_spharm_vae package</a><input class="toctree-checkbox" id="toctree-checkbox-24" name="toctree-checkbox-24" role="switch" type="checkbox"/><label for="toctree-checkbox-24"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.models.vae.o2_spharm_vae.o2_spharm_encoder.html">cyto_dl.models.vae.o2_spharm_vae.o2_spharm_encoder module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.models.vae.o2_spharm_vae.o2_spharm_vae.html">cyto_dl.models.vae.o2_spharm_vae.o2_spharm_vae module</a></li>
+</ul>
+</li>
+<li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.models.vae.priors.html">cyto_dl.models.vae.priors package</a><input class="toctree-checkbox" id="toctree-checkbox-25" name="toctree-checkbox-25" role="switch" type="checkbox"/><label for="toctree-checkbox-25"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.models.vae.priors.abstract_prior.html">cyto_dl.models.vae.priors.abstract_prior module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.models.vae.priors.gaussian.html">cyto_dl.models.vae.priors.gaussian module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.models.vae.priors.identity_prior.html">cyto_dl.models.vae.priors.identity_prior module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.models.vae.priors.joint_prior.html">cyto_dl.models.vae.priors.joint_prior module</a></li>
+</ul>
+</li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.models.vae.base_vae.html">cyto_dl.models.vae.base_vae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.models.vae.image_canon_vae.html">cyto_dl.models.vae.image_canon_vae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.models.vae.image_encoder.html">cyto_dl.models.vae.image_encoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.models.vae.image_vae.html">cyto_dl.models.vae.image_vae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.models.vae.latent_loss_vae.html">cyto_dl.models.vae.latent_loss_vae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.models.vae.point_cloud_vae.html">cyto_dl.models.vae.point_cloud_vae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.models.vae.tabular_vae.html">cyto_dl.models.vae.tabular_vae module</a></li>
+</ul>
+</li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.models.base_model.html">cyto_dl.models.base_model module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.models.basic_model.html">cyto_dl.models.basic_model module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.html">cyto_dl.nn package</a><input class="toctree-checkbox" id="toctree-checkbox-26" name="toctree-checkbox-26" role="switch" type="checkbox"/><label for="toctree-checkbox-26"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.discriminators.html">cyto_dl.nn.discriminators package</a><input class="toctree-checkbox" id="toctree-checkbox-27" name="toctree-checkbox-27" role="switch" type="checkbox"/><label for="toctree-checkbox-27"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.discriminators.multi_scale_discriminator.html">cyto_dl.nn.discriminators.multi_scale_discriminator module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.discriminators.n_layer_discriminator.html">cyto_dl.nn.discriminators.n_layer_discriminator module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.head.html">cyto_dl.nn.head package</a><input class="toctree-checkbox" id="toctree-checkbox-28" name="toctree-checkbox-28" role="switch" type="checkbox"/><label for="toctree-checkbox-28"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.head.base_head.html">cyto_dl.nn.head.base_head module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.head.gan_head.html">cyto_dl.nn.head.gan_head module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.head.gan_head_superres.html">cyto_dl.nn.head.gan_head_superres module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.head.mae_head.html">cyto_dl.nn.head.mae_head module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.head.mask_head.html">cyto_dl.nn.head.mask_head module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.head.res_blocks_head.html">cyto_dl.nn.head.res_blocks_head module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.head.vic_reg.html">cyto_dl.nn.head.vic_reg module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.losses.html">cyto_dl.nn.losses package</a><input class="toctree-checkbox" id="toctree-checkbox-29" name="toctree-checkbox-29" role="switch" type="checkbox"/><label for="toctree-checkbox-29"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.losses.adversarial_loss.html">cyto_dl.nn.losses.adversarial_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.losses.chamfer_loss.html">cyto_dl.nn.losses.chamfer_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.losses.continuous_bernoulli.html">cyto_dl.nn.losses.continuous_bernoulli module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.losses.cosine_loss.html">cyto_dl.nn.losses.cosine_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.losses.gan_loss.html">cyto_dl.nn.losses.gan_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.losses.gaussian_nll_loss.html">cyto_dl.nn.losses.gaussian_nll_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.losses.geomloss.html">cyto_dl.nn.losses.geomloss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.losses.loss_wrapper.html">cyto_dl.nn.losses.loss_wrapper module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.losses.spharm_loss.html">cyto_dl.nn.losses.spharm_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.losses.threshold_loss.html">cyto_dl.nn.losses.threshold_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.losses.vic_reg.html">cyto_dl.nn.losses.vic_reg module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.losses.weibull.html">cyto_dl.nn.losses.weibull module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.losses.weighted_mse_loss.html">cyto_dl.nn.losses.weighted_mse_loss module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.point_cloud.html">cyto_dl.nn.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-30" name="toctree-checkbox-30" role="switch" type="checkbox"/><label for="toctree-checkbox-30"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.point_cloud.dgcnn.html">cyto_dl.nn.point_cloud.dgcnn module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.point_cloud.folding_net.html">cyto_dl.nn.point_cloud.folding_net module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.point_cloud.graph_functions.html">cyto_dl.nn.point_cloud.graph_functions module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.point_cloud.vnn.html">cyto_dl.nn.point_cloud.vnn module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
+</ul>
+</li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
+</ul>
+</li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.hr_skip.html">cyto_dl.nn.hr_skip module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.mlp.html">cyto_dl.nn.mlp module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.res_unit.html">cyto_dl.nn.res_unit module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.spatial_transformer.html">cyto_dl.nn.spatial_transformer module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
+</ul>
+</li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
+</ul>
+</li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.utils.array.html">cyto_dl.utils.array module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.utils.checkpoint.html">cyto_dl.utils.checkpoint module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.utils.config.html">cyto_dl.utils.config module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.utils.download_test_data.html">cyto_dl.utils.download_test_data module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.utils.dummy_dataset.html">cyto_dl.utils.dummy_dataset module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.utils.pylogger.html">cyto_dl.utils.pylogger module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.utils.rich_utils.html">cyto_dl.utils.rich_utils module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.utils.rotation.html">cyto_dl.utils.rotation module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.utils.template_utils.html">cyto_dl.utils.template_utils module</a></li>
+</ul>
+</li>
+<li class="toctree-l3"><a class="reference internal" href="../../../../cyto_dl.compile.html">cyto_dl.compile module</a></li>
+<li class="toctree-l3"><a class="reference internal" href="../../../../cyto_dl.eval.html">cyto_dl.eval module</a></li>
+<li class="toctree-l3"><a class="reference internal" href="../../../../cyto_dl.train.html">cyto_dl.train module</a></li>
+</ul>
+</li>
+</ul>
+</li>
+</ul>
+
+</div>
+</div>
+
+      </div>
+      
+    </div>
+  </aside>
+  <div class="main">
+    <div class="content">
+      <div class="article-container">
+        <a href="#" class="back-to-top muted-link">
+          <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24">
+            <path d="M13 20h-2V8l-5.5 5.5-1.42-1.42L12 4.16l7.92 7.92-1.42 1.42L13 8v12z"></path>
+          </svg>
+          <span>Back to top</span>
+        </a>
+        <div class="content-icon-container">
+          <div class="theme-toggle-container theme-toggle-content">
+            <button class="theme-toggle">
+              <div class="visually-hidden">Toggle Light / Dark / Auto color theme</div>
+              <svg class="theme-icon-when-auto"><use href="#svg-sun-half"></use></svg>
+              <svg class="theme-icon-when-dark"><use href="#svg-moon"></use></svg>
+              <svg class="theme-icon-when-light"><use href="#svg-sun"></use></svg>
+            </button>
+          </div>
+          <label class="toc-overlay-icon toc-content-icon no-toc" for="__toc">
+            <div class="visually-hidden">Toggle table of contents sidebar</div>
+            <i class="icon"><svg><use href="#svg-toc"></use></svg></i>
+          </label>
+        </div>
+        <article role="main">
+          <h1>Source code for cyto_dl.nn.vits.encoder</h1><div class="highlight"><pre>
+<span></span><span class="c1"># modified from https://github.com/IcarusWizard/MAE/blob/main/model.py#L124</span>
+<span class="c1"># inspired by https://github.com/facebookresearch/hiera</span>
+
+<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Dict</span><span class="p">,</span> <span class="n">List</span><span class="p">,</span> <span class="n">Optional</span><span class="p">,</span> <span class="n">Union</span>
+
+<span class="kn">import</span> <span class="nn">numpy</span> <span class="k">as</span> <span class="nn">np</span>
+<span class="kn">import</span> <span class="nn">torch</span>
+<span class="kn">import</span> <span class="nn">torch.nn</span> <span class="k">as</span> <span class="nn">nn</span>
+<span class="kn">import</span> <span class="nn">torch.nn.functional</span>
+<span class="kn">from</span> <span class="nn">einops</span> <span class="kn">import</span> <span class="n">rearrange</span>
+<span class="kn">from</span> <span class="nn">einops.layers.torch</span> <span class="kn">import</span> <span class="n">Rearrange</span>
+<span class="kn">from</span> <span class="nn">timm.models.layers</span> <span class="kn">import</span> <span class="n">trunc_normal_</span>
+<span class="kn">from</span> <span class="nn">timm.models.vision_transformer</span> <span class="kn">import</span> <span class="n">Block</span>
+
+<span class="kn">from</span> <span class="nn">cyto_dl.nn.vits.blocks</span> <span class="kn">import</span> <span class="n">IntermediateWeigher</span><span class="p">,</span> <span class="n">Patchify</span>
+<span class="kn">from</span> <span class="nn">cyto_dl.nn.vits.blocks.masked_unit_attention</span> <span class="kn">import</span> <span class="n">HieraBlock</span>
+<span class="kn">from</span> <span class="nn">cyto_dl.nn.vits.blocks.patchify</span> <span class="kn">import</span> <span class="n">PatchifyHiera</span>
+<span class="kn">from</span> <span class="nn">cyto_dl.nn.vits.utils</span> <span class="kn">import</span> <span class="n">match_tuple_dimensions</span>
+
+
+<div class="viewcode-block" id="MAE_Encoder"><a class="viewcode-back" href="../../../../cyto_dl.nn.vits.encoder.html#cyto_dl.nn.vits.encoder.MAE_Encoder">[docs]</a><span class="k">class</span> <span class="nc">MAE_Encoder</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
+    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span>
+        <span class="bp">self</span><span class="p">,</span>
+        <span class="n">num_patches</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span>
+        <span class="n">spatial_dims</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">3</span><span class="p">,</span>
+        <span class="n">patch_size</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">]]</span> <span class="o">=</span> <span class="mi">4</span><span class="p">,</span>
+        <span class="n">emb_dim</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">192</span><span class="p">,</span>
+        <span class="n">num_layer</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">12</span><span class="p">,</span>
+        <span class="n">num_head</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">3</span><span class="p">,</span>
+        <span class="n">context_pixels</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">]]]</span> <span class="o">=</span> <span class="mi">0</span><span class="p">,</span>
+        <span class="n">input_channels</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">1</span><span class="p">,</span>
+        <span class="n">n_intermediate_weights</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="o">-</span><span class="mi">1</span><span class="p">,</span>
+    <span class="p">)</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
+<span class="w">        </span><span class="sd">&quot;&quot;&quot;</span>
+<span class="sd">        Parameters</span>
+<span class="sd">        ----------</span>
+<span class="sd">        num_patches: List[int], int</span>
+<span class="sd">            Number of patches in each dimension. If a single int is provided, the number of patches in each dimension will be the same.</span>
+<span class="sd">        spatial_dims: int</span>
+<span class="sd">            Number of spatial dimensions</span>
+<span class="sd">        patch_size: List[int]</span>
+<span class="sd">            Size of each patch</span>
+<span class="sd">        emb_dim: int</span>
+<span class="sd">            Dimension of embedding</span>
+<span class="sd">        num_layer: int</span>
+<span class="sd">            Number of transformer layers</span>
+<span class="sd">        num_head: int</span>
+<span class="sd">            Number of heads in transformer</span>
+<span class="sd">        context_pixels: List[int], int</span>
+<span class="sd">            Number of extra pixels around each patch to include in convolutional embedding to encoder dimension. If a single int is provided, the number of context pixels in each dimension will be the same.</span>
+<span class="sd">        input_channels: int</span>
+<span class="sd">            Number of input channels</span>
+<span class="sd">        n_intermediate_weights: bool</span>
+<span class="sd">            Whether to use intermediate weights for weighted sum of intermediate layers</span>
+<span class="sd">        &quot;&quot;&quot;</span>
+        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>
+        <span class="n">num_patches</span><span class="p">,</span> <span class="n">patch_size</span><span class="p">,</span> <span class="n">context_pixels</span> <span class="o">=</span> <span class="n">match_tuple_dimensions</span><span class="p">(</span>
+            <span class="n">spatial_dims</span><span class="p">,</span> <span class="p">[</span><span class="n">num_patches</span><span class="p">,</span> <span class="n">patch_size</span><span class="p">,</span> <span class="n">context_pixels</span><span class="p">]</span>
+        <span class="p">)</span>
+
+        <span class="bp">self</span><span class="o">.</span><span class="n">cls_token</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Parameter</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="n">emb_dim</span><span class="p">))</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">patchify</span> <span class="o">=</span> <span class="n">Patchify</span><span class="p">(</span>
+            <span class="n">patch_size</span><span class="p">,</span> <span class="n">emb_dim</span><span class="p">,</span> <span class="n">num_patches</span><span class="p">,</span> <span class="n">spatial_dims</span><span class="p">,</span> <span class="n">context_pixels</span><span class="p">,</span> <span class="n">input_channels</span>
+        <span class="p">)</span>
+        <span class="n">weight_intermediates</span> <span class="o">=</span> <span class="n">n_intermediate_weights</span> <span class="o">&gt;</span> <span class="mi">0</span>
+        <span class="k">if</span> <span class="n">weight_intermediates</span><span class="p">:</span>
+            <span class="bp">self</span><span class="o">.</span><span class="n">transformer</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">ModuleList</span><span class="p">(</span>
+                <span class="p">[</span><span class="n">Block</span><span class="p">(</span><span class="n">emb_dim</span><span class="p">,</span> <span class="n">num_head</span><span class="p">)</span> <span class="k">for</span> <span class="n">_</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">num_layer</span><span class="p">)]</span>
+            <span class="p">)</span>
+        <span class="k">else</span><span class="p">:</span>
+            <span class="bp">self</span><span class="o">.</span><span class="n">transformer</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Sequential</span><span class="p">(</span>
+                <span class="o">*</span><span class="p">[</span><span class="n">Block</span><span class="p">(</span><span class="n">emb_dim</span><span class="p">,</span> <span class="n">num_head</span><span class="p">)</span> <span class="k">for</span> <span class="n">_</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">num_layer</span><span class="p">)]</span>
+            <span class="p">)</span>
+
+        <span class="bp">self</span><span class="o">.</span><span class="n">layer_norm</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">LayerNorm</span><span class="p">(</span><span class="n">emb_dim</span><span class="p">)</span>
+
+        <span class="bp">self</span><span class="o">.</span><span class="n">intermediate_weighter</span> <span class="o">=</span> <span class="p">(</span>
+            <span class="n">IntermediateWeigher</span><span class="p">(</span><span class="n">num_layer</span><span class="p">,</span> <span class="n">emb_dim</span><span class="p">,</span> <span class="n">n_intermediate_weights</span><span class="p">)</span>
+            <span class="k">if</span> <span class="n">weight_intermediates</span>
+            <span class="k">else</span> <span class="kc">None</span>
+        <span class="p">)</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">init_weight</span><span class="p">()</span>
+
+<div class="viewcode-block" id="MAE_Encoder.init_weight"><a class="viewcode-back" href="../../../../cyto_dl.nn.vits.encoder.html#cyto_dl.nn.vits.encoder.MAE_Encoder.init_weight">[docs]</a>    <span class="k">def</span> <span class="nf">init_weight</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
+        <span class="n">trunc_normal_</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">cls_token</span><span class="p">,</span> <span class="n">std</span><span class="o">=</span><span class="mf">0.02</span><span class="p">)</span></div>
+
+<div class="viewcode-block" id="MAE_Encoder.forward"><a class="viewcode-back" href="../../../../cyto_dl.nn.vits.encoder.html#cyto_dl.nn.vits.encoder.MAE_Encoder.forward">[docs]</a>    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">img</span><span class="p">,</span> <span class="n">mask_ratio</span><span class="o">=</span><span class="mf">0.75</span><span class="p">):</span>
+        <span class="n">patches</span><span class="p">,</span> <span class="n">mask</span><span class="p">,</span> <span class="n">forward_indexes</span><span class="p">,</span> <span class="n">backward_indexes</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">patchify</span><span class="p">(</span><span class="n">img</span><span class="p">,</span> <span class="n">mask_ratio</span><span class="p">)</span>
+        <span class="n">patches</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">cat</span><span class="p">([</span><span class="bp">self</span><span class="o">.</span><span class="n">cls_token</span><span class="o">.</span><span class="n">expand</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="n">patches</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class="o">-</span><span class="mi">1</span><span class="p">),</span> <span class="n">patches</span><span class="p">],</span> <span class="n">dim</span><span class="o">=</span><span class="mi">0</span><span class="p">)</span>
+        <span class="n">patches</span> <span class="o">=</span> <span class="n">rearrange</span><span class="p">(</span><span class="n">patches</span><span class="p">,</span> <span class="s2">&quot;t b c -&gt; b t c&quot;</span><span class="p">)</span>
+
+        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">intermediate_weighter</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
+            <span class="n">intermediates</span> <span class="o">=</span> <span class="p">[</span><span class="n">patches</span><span class="p">]</span>
+            <span class="k">for</span> <span class="n">block</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">transformer</span><span class="p">:</span>
+                <span class="n">patches</span> <span class="o">=</span> <span class="n">block</span><span class="p">(</span><span class="n">patches</span><span class="p">)</span>
+                <span class="n">intermediates</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">patches</span><span class="p">)</span>
+            <span class="n">features</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">layer_norm</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">intermediate_weighter</span><span class="p">(</span><span class="n">intermediates</span><span class="p">))</span>
+            <span class="n">features</span> <span class="o">=</span> <span class="n">rearrange</span><span class="p">(</span><span class="n">features</span><span class="p">,</span> <span class="s2">&quot;n b t c -&gt; n t b c&quot;</span><span class="p">)</span>
+        <span class="k">else</span><span class="p">:</span>
+            <span class="n">features</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">layer_norm</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">transformer</span><span class="p">(</span><span class="n">patches</span><span class="p">))</span>
+            <span class="n">features</span> <span class="o">=</span> <span class="n">rearrange</span><span class="p">(</span><span class="n">features</span><span class="p">,</span> <span class="s2">&quot;b t c -&gt; t b c&quot;</span><span class="p">)</span>
+        <span class="k">if</span> <span class="n">mask_ratio</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>
+            <span class="k">return</span> <span class="n">features</span><span class="p">,</span> <span class="n">mask</span><span class="p">,</span> <span class="n">forward_indexes</span><span class="p">,</span> <span class="n">backward_indexes</span>
+        <span class="k">return</span> <span class="n">features</span></div></div>
+
+
+<div class="viewcode-block" id="JEPAEncoder"><a class="viewcode-back" href="../../../../cyto_dl.nn.vits.encoder.html#cyto_dl.nn.vits.encoder.JEPAEncoder">[docs]</a><span class="k">class</span> <span class="nc">JEPAEncoder</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
+    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span>
+        <span class="bp">self</span><span class="p">,</span>
+        <span class="n">num_patches</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">]],</span>
+        <span class="n">spatial_dims</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">3</span><span class="p">,</span>
+        <span class="n">patch_size</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">]]</span> <span class="o">=</span> <span class="mi">4</span><span class="p">,</span>
+        <span class="n">emb_dim</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">192</span><span class="p">,</span>
+        <span class="n">num_layer</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">12</span><span class="p">,</span>
+        <span class="n">num_head</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">3</span><span class="p">,</span>
+        <span class="n">context_pixels</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">]]]</span> <span class="o">=</span> <span class="mi">0</span><span class="p">,</span>
+        <span class="n">input_channels</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">1</span><span class="p">,</span>
+        <span class="n">learnable_pos_embedding</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">True</span><span class="p">,</span>
+    <span class="p">)</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
+<span class="w">        </span><span class="sd">&quot;&quot;&quot;</span>
+<span class="sd">        Parameters</span>
+<span class="sd">        ----------</span>
+<span class="sd">        num_patches: List[int], int</span>
+<span class="sd">            Number of patches in each dimension. If a single int is provided, the number of patches in each dimension will be the same.</span>
+<span class="sd">        spatial_dims: int</span>
+<span class="sd">            Number of spatial dimensions</span>
+<span class="sd">        patch_size: List[int]</span>
+<span class="sd">            Size of each patch</span>
+<span class="sd">        emb_dim: int</span>
+<span class="sd">            Dimension of embedding</span>
+<span class="sd">        num_layer: int</span>
+<span class="sd">            Number of transformer layers</span>
+<span class="sd">        num_head: int</span>
+<span class="sd">            Number of heads in transformer</span>
+<span class="sd">        context_pixels: List[int], int</span>
+<span class="sd">            Number of extra pixels around each patch to include in convolutional embedding to encoder dimension. If a single int is provided, the number of context pixels in each dimension will be the same.</span>
+<span class="sd">        input_channels: int</span>
+<span class="sd">            Number of input channels</span>
+<span class="sd">        learnable_pos_embedding: bool</span>
+<span class="sd">            If True, learnable positional embeddings are used. If False, fixed sin/cos positional embeddings. Empirically, fixed positional embeddings work better for brightfield images.</span>
+<span class="sd">        &quot;&quot;&quot;</span>
+        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>
+        <span class="n">num_patches</span><span class="p">,</span> <span class="n">patch_size</span><span class="p">,</span> <span class="n">context_pixels</span> <span class="o">=</span> <span class="n">match_tuple_dimensions</span><span class="p">(</span>
+            <span class="n">spatial_dims</span><span class="p">,</span> <span class="p">[</span><span class="n">num_patches</span><span class="p">,</span> <span class="n">patch_size</span><span class="p">,</span> <span class="n">context_pixels</span><span class="p">]</span>
+        <span class="p">)</span>
+
+        <span class="bp">self</span><span class="o">.</span><span class="n">patchify</span> <span class="o">=</span> <span class="n">Patchify</span><span class="p">(</span>
+            <span class="n">patch_size</span><span class="o">=</span><span class="n">patch_size</span><span class="p">,</span>
+            <span class="n">emb_dim</span><span class="o">=</span><span class="n">emb_dim</span><span class="p">,</span>
+            <span class="n">n_patches</span><span class="o">=</span><span class="n">num_patches</span><span class="p">,</span>
+            <span class="n">spatial_dims</span><span class="o">=</span><span class="n">spatial_dims</span><span class="p">,</span>
+            <span class="n">context_pixels</span><span class="o">=</span><span class="n">context_pixels</span><span class="p">,</span>
+            <span class="n">input_channels</span><span class="o">=</span><span class="n">input_channels</span><span class="p">,</span>
+            <span class="n">learnable_pos_embedding</span><span class="o">=</span><span class="n">learnable_pos_embedding</span><span class="p">,</span>
+        <span class="p">)</span>
+
+        <span class="bp">self</span><span class="o">.</span><span class="n">transformer</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Sequential</span><span class="p">(</span>
+            <span class="o">*</span><span class="p">[</span><span class="n">Block</span><span class="p">(</span><span class="n">emb_dim</span><span class="p">,</span> <span class="n">num_head</span><span class="p">)</span> <span class="k">for</span> <span class="n">_</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">num_layer</span><span class="p">)]</span>
+        <span class="p">)</span>
+
+        <span class="bp">self</span><span class="o">.</span><span class="n">layer_norm</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">LayerNorm</span><span class="p">(</span><span class="n">emb_dim</span><span class="p">)</span>
+
+<div class="viewcode-block" id="JEPAEncoder.forward"><a class="viewcode-back" href="../../../../cyto_dl.nn.vits.encoder.html#cyto_dl.nn.vits.encoder.JEPAEncoder.forward">[docs]</a>    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">img</span><span class="p">,</span> <span class="n">patchify</span><span class="o">=</span><span class="kc">True</span><span class="p">):</span>
+        <span class="k">if</span> <span class="n">patchify</span><span class="p">:</span>
+            <span class="n">patches</span><span class="p">,</span> <span class="n">_</span><span class="p">,</span> <span class="n">_</span><span class="p">,</span> <span class="n">_</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">patchify</span><span class="p">(</span><span class="n">img</span><span class="p">,</span> <span class="n">mask_ratio</span><span class="o">=</span><span class="mi">0</span><span class="p">)</span>
+        <span class="k">else</span><span class="p">:</span>
+            <span class="n">patches</span> <span class="o">=</span> <span class="n">img</span>
+        <span class="n">patches</span> <span class="o">=</span> <span class="n">rearrange</span><span class="p">(</span><span class="n">patches</span><span class="p">,</span> <span class="s2">&quot;t b c -&gt; b t c&quot;</span><span class="p">)</span>
+        <span class="n">features</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">layer_norm</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">transformer</span><span class="p">(</span><span class="n">patches</span><span class="p">))</span>
+        <span class="k">return</span> <span class="n">features</span></div></div>
+
+
+<div class="viewcode-block" id="SpatialMerger"><a class="viewcode-back" href="../../../../cyto_dl.nn.vits.encoder.html#cyto_dl.nn.vits.encoder.SpatialMerger">[docs]</a><span class="k">class</span> <span class="nc">SpatialMerger</span><span class="p">(</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
+<span class="w">    </span><span class="sd">&quot;&quot;&quot;Class for converting multi-resolution Hiera features to the same (lowest) spatial resolution</span>
+<span class="sd">    via convolution.&quot;&quot;&quot;</span>
+
+    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span>
+        <span class="bp">self</span><span class="p">,</span> <span class="n">downsample_factor</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span> <span class="n">in_dim</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">out_dim</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">spatial_dims</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">3</span>
+    <span class="p">):</span>
+        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>
+        <span class="n">downsample_factor</span> <span class="o">=</span> <span class="n">match_tuple_dimensions</span><span class="p">(</span><span class="n">spatial_dims</span><span class="p">,</span> <span class="p">[</span><span class="n">downsample_factor</span><span class="p">])[</span><span class="mi">0</span><span class="p">]</span>
+
+        <span class="bp">self</span><span class="o">.</span><span class="n">spatial_dims</span> <span class="o">=</span> <span class="n">spatial_dims</span>
+        <span class="n">conv_fn</span> <span class="o">=</span> <span class="n">nn</span><span class="o">.</span><span class="n">Conv3d</span> <span class="k">if</span> <span class="n">spatial_dims</span> <span class="o">==</span> <span class="mi">3</span> <span class="k">else</span> <span class="n">nn</span><span class="o">.</span><span class="n">Conv2d</span>
+        <span class="n">conv</span> <span class="o">=</span> <span class="n">conv_fn</span><span class="p">(</span>
+            <span class="n">in_channels</span><span class="o">=</span><span class="n">in_dim</span><span class="p">,</span>
+            <span class="n">out_channels</span><span class="o">=</span><span class="n">out_dim</span><span class="p">,</span>
+            <span class="n">kernel_size</span><span class="o">=</span><span class="n">downsample_factor</span><span class="p">,</span>
+            <span class="n">stride</span><span class="o">=</span><span class="n">downsample_factor</span><span class="p">,</span>
+            <span class="n">padding</span><span class="o">=</span><span class="mi">0</span><span class="p">,</span>
+            <span class="n">bias</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
+        <span class="p">)</span>
+        <span class="k">if</span> <span class="n">spatial_dims</span> <span class="o">==</span> <span class="mi">3</span><span class="p">:</span>
+            <span class="n">tokens2img</span> <span class="o">=</span> <span class="n">Rearrange</span><span class="p">(</span>
+                <span class="s2">&quot;b n_mu (z y x) c -&gt; (b n_mu) c z y x&quot;</span><span class="p">,</span>
+                <span class="n">z</span><span class="o">=</span><span class="n">downsample_factor</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span>
+                <span class="n">y</span><span class="o">=</span><span class="n">downsample_factor</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span>
+                <span class="n">x</span><span class="o">=</span><span class="n">downsample_factor</span><span class="p">[</span><span class="mi">2</span><span class="p">],</span>
+            <span class="p">)</span>
+        <span class="k">else</span><span class="p">:</span>
+            <span class="n">tokens2img</span> <span class="o">=</span> <span class="n">Rearrange</span><span class="p">(</span>
+                <span class="s2">&quot;b n_mu (y x) c -&gt; (b n_mu) c y x&quot;</span><span class="p">,</span>
+                <span class="n">y</span><span class="o">=</span><span class="n">downsample_factor</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span>
+                <span class="n">x</span><span class="o">=</span><span class="n">downsample_factor</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span>
+            <span class="p">)</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">model</span> <span class="o">=</span> <span class="n">nn</span><span class="o">.</span><span class="n">Sequential</span><span class="p">(</span><span class="n">tokens2img</span><span class="p">,</span> <span class="n">conv</span><span class="p">)</span>
+
+<div class="viewcode-block" id="SpatialMerger.forward"><a class="viewcode-back" href="../../../../cyto_dl.nn.vits.encoder.html#cyto_dl.nn.vits.encoder.SpatialMerger.forward">[docs]</a>    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x</span><span class="p">):</span>
+        <span class="n">b</span><span class="p">,</span> <span class="n">n_mu</span><span class="p">,</span> <span class="n">_</span><span class="p">,</span> <span class="n">_</span> <span class="o">=</span> <span class="n">x</span><span class="o">.</span><span class="n">shape</span>
+        <span class="n">x</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">model</span><span class="p">(</span><span class="n">x</span><span class="p">)</span>
+        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">spatial_dims</span> <span class="o">==</span> <span class="mi">3</span><span class="p">:</span>
+            <span class="n">x</span> <span class="o">=</span> <span class="n">rearrange</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="s2">&quot;(b n_mu) c z y x -&gt; b n_mu (z y x) c&quot;</span><span class="p">,</span> <span class="n">b</span><span class="o">=</span><span class="n">b</span><span class="p">,</span> <span class="n">n_mu</span><span class="o">=</span><span class="n">n_mu</span><span class="p">)</span>
+        <span class="k">else</span><span class="p">:</span>
+            <span class="n">x</span> <span class="o">=</span> <span class="n">rearrange</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="s2">&quot;(b n_mu) c y x -&gt; b n_mu (y x) c&quot;</span><span class="p">,</span> <span class="n">b</span><span class="o">=</span><span class="n">b</span><span class="p">,</span> <span class="n">n_mu</span><span class="o">=</span><span class="n">n_mu</span><span class="p">)</span>
+        <span class="k">return</span> <span class="n">x</span></div></div>
+
+
+<div class="viewcode-block" id="HieraEncoder"><a class="viewcode-back" href="../../../../cyto_dl.nn.vits.encoder.html#cyto_dl.nn.vits.encoder.HieraEncoder">[docs]</a><span class="k">class</span> <span class="nc">HieraEncoder</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
+    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span>
+        <span class="bp">self</span><span class="p">,</span>
+        <span class="n">num_patches</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">]],</span>
+        <span class="n">num_mask_units</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">]],</span>
+        <span class="n">architecture</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="n">Dict</span><span class="p">],</span>
+        <span class="n">emb_dim</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">64</span><span class="p">,</span>
+        <span class="n">spatial_dims</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">3</span><span class="p">,</span>
+        <span class="n">patch_size</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">]]</span> <span class="o">=</span> <span class="mi">4</span><span class="p">,</span>
+        <span class="n">context_pixels</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">]]]</span> <span class="o">=</span> <span class="mi">0</span><span class="p">,</span>
+        <span class="n">input_channels</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">1</span><span class="p">,</span>
+        <span class="n">save_layers</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span>
+    <span class="p">)</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
+<span class="w">        </span><span class="sd">&quot;&quot;&quot;</span>
+<span class="sd">        Parameters</span>
+<span class="sd">        ----------</span>
+<span class="sd">        num_patches: int, List[int]</span>
+<span class="sd">            Number of patches in each dimension. If a single int is provided, the number of patches in each dimension will be the same.</span>
+<span class="sd">        num_mask_units: int, List[int]</span>
+<span class="sd">            Number of mask units in each dimension. If a single int is provided, the number of mask units in each dimension will be the same.</span>
+<span class="sd">        architecture: List[Dict]</span>
+<span class="sd">            List of dictionaries specifying the architecture of the transformer. Each dictionary should have the following keys:</span>
+<span class="sd">            - repeat: int</span>
+<span class="sd">                Number of times to repeat the block</span>
+<span class="sd">            - num_heads: int</span>
+<span class="sd">                Number of heads in the multihead attention</span>
+<span class="sd">            - q_stride: int, List[int]</span>
+<span class="sd">                Stride for the query in each spatial dimension</span>
+<span class="sd">            - self_attention: bool</span>
+<span class="sd">                Whether to use self attention or mask unit attention</span>
+<span class="sd">            On the last repeat of each non-self-attention block, the embedding dimension is doubled and spatial pooling with `q_stride` is performed within each mask unit. For example, a block with a embed_dim=4, q_stride=2, and repeat=2, the first repeat just does mask unit attention, while the second will produce an 8-dimensional output that has been spatially pooled.</span>
+<span class="sd">        emb_dim: int</span>
+<span class="sd">            Dimension of embedding</span>
+<span class="sd">        spatial_dims: int</span>
+<span class="sd">            Number of spatial dimensions</span>
+<span class="sd">        patch_size: List[int]</span>
+<span class="sd">            Size of each patch</span>
+<span class="sd">        context_pixels: List[int]</span>
+<span class="sd">            Number of extra pixels around each patch to include in convolutional embedding to encoder dimension.</span>
+<span class="sd">        input_channels: int</span>
+<span class="sd">            Number of input channels</span>
+<span class="sd">        save_layers: bool</span>
+<span class="sd">            Whether to save the intermediate layer outputs</span>
+<span class="sd">        &quot;&quot;&quot;</span>
+        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>
+        <span class="n">num_patches</span><span class="p">,</span> <span class="n">num_mask_units</span><span class="p">,</span> <span class="n">patch_size</span><span class="p">,</span> <span class="n">context_pixels</span> <span class="o">=</span> <span class="n">match_tuple_dimensions</span><span class="p">(</span>
+            <span class="n">spatial_dims</span><span class="p">,</span> <span class="p">[</span><span class="n">num_patches</span><span class="p">,</span> <span class="n">num_mask_units</span><span class="p">,</span> <span class="n">patch_size</span><span class="p">,</span> <span class="n">context_pixels</span><span class="p">]</span>
+        <span class="p">)</span>
+        <span class="c1"># make sure q stride shape matches spatial dims</span>
+        <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">architecture</span><span class="p">)):</span>
+            <span class="k">if</span> <span class="s2">&quot;q_stride&quot;</span> <span class="ow">in</span> <span class="n">architecture</span><span class="p">[</span><span class="n">i</span><span class="p">]:</span>
+                <span class="n">architecture</span><span class="p">[</span><span class="n">i</span><span class="p">][</span><span class="s2">&quot;q_stride&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">match_tuple_dimensions</span><span class="p">(</span>
+                    <span class="n">spatial_dims</span><span class="p">,</span> <span class="p">[</span><span class="n">architecture</span><span class="p">[</span><span class="n">i</span><span class="p">][</span><span class="s2">&quot;q_stride&quot;</span><span class="p">]]</span>
+                <span class="p">)[</span><span class="mi">0</span><span class="p">]</span>
+
+        <span class="bp">self</span><span class="o">.</span><span class="n">save_layers</span> <span class="o">=</span> <span class="n">save_layers</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">patchify</span> <span class="o">=</span> <span class="n">PatchifyHiera</span><span class="p">(</span>
+            <span class="n">patch_size</span><span class="o">=</span><span class="n">patch_size</span><span class="p">,</span>
+            <span class="n">n_patches</span><span class="o">=</span><span class="n">num_patches</span><span class="p">,</span>
+            <span class="n">emb_dim</span><span class="o">=</span><span class="n">emb_dim</span><span class="p">,</span>
+            <span class="n">spatial_dims</span><span class="o">=</span><span class="n">spatial_dims</span><span class="p">,</span>
+            <span class="n">context_pixels</span><span class="o">=</span><span class="n">context_pixels</span><span class="p">,</span>
+            <span class="n">input_channels</span><span class="o">=</span><span class="n">input_channels</span><span class="p">,</span>
+            <span class="n">mask_units_per_dim</span><span class="o">=</span><span class="n">num_mask_units</span><span class="p">,</span>
+        <span class="p">)</span>
+
+        <span class="n">patches_per_mask_unit</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">(</span><span class="n">num_patches</span><span class="p">)</span> <span class="o">//</span> <span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">(</span><span class="n">num_mask_units</span><span class="p">)</span>
+
+        <span class="n">total_downsampling_per_axis</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">prod</span><span class="p">(</span>
+            <span class="p">[</span><span class="n">block</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;q_stride&quot;</span><span class="p">,</span> <span class="p">[</span><span class="mi">1</span><span class="p">]</span> <span class="o">*</span> <span class="n">spatial_dims</span><span class="p">)</span> <span class="k">for</span> <span class="n">block</span> <span class="ow">in</span> <span class="n">architecture</span><span class="p">],</span> <span class="n">axis</span><span class="o">=</span><span class="mi">0</span>
+        <span class="p">)</span>
+
+        <span class="k">assert</span> <span class="n">np</span><span class="o">.</span><span class="n">all</span><span class="p">(</span>
+            <span class="n">patches_per_mask_unit</span> <span class="o">-</span> <span class="n">total_downsampling_per_axis</span> <span class="o">&gt;=</span> <span class="mi">0</span>
+        <span class="p">),</span> <span class="sa">f</span><span class="s2">&quot;Number of mask units must be greater than the total downsampling ratio, got </span><span class="si">{</span><span class="n">patches_per_mask_unit</span><span class="si">}</span><span class="s2"> patches per mask unit and </span><span class="si">{</span><span class="n">total_downsampling_per_axis</span><span class="si">}</span><span class="s2"> total downsampling ratio. Please adjust your q_stride or increase the number of patches per mask unit.&quot;</span>
+        <span class="k">assert</span> <span class="n">np</span><span class="o">.</span><span class="n">all</span><span class="p">(</span>
+            <span class="n">patches_per_mask_unit</span> <span class="o">%</span> <span class="n">total_downsampling_per_axis</span> <span class="o">==</span> <span class="mi">0</span>
+        <span class="p">),</span> <span class="sa">f</span><span class="s2">&quot;Number of mask units must be divisible by the total downsampling ratio, got </span><span class="si">{</span><span class="n">patches_per_mask_unit</span><span class="si">}</span><span class="s2"> patches per mask unit and </span><span class="si">{</span><span class="n">total_downsampling_per_axis</span><span class="si">}</span><span class="s2"> total downsampling ratio. Please adjust your q_stride&quot;</span>
+
+        <span class="c1"># number of output features doubles in each masked unit attention block, stays constant during self attention blocks</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">final_dim</span> <span class="o">=</span> <span class="n">emb_dim</span> <span class="o">*</span> <span class="p">(</span><span class="mi">2</span> <span class="o">**</span> <span class="nb">len</span><span class="p">(</span><span class="n">architecture</span><span class="p">))</span>
+
+        <span class="bp">self</span><span class="o">.</span><span class="n">save_block_idxs</span> <span class="o">=</span> <span class="p">[]</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">save_block_dims</span> <span class="o">=</span> <span class="p">[]</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">spatial_mergers</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">ParameterDict</span><span class="p">({})</span>
+        <span class="n">transformer</span> <span class="o">=</span> <span class="p">[]</span>
+        <span class="n">num_blocks</span> <span class="o">=</span> <span class="mi">0</span>
+        <span class="k">for</span> <span class="n">stage_num</span><span class="p">,</span> <span class="n">stage</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">architecture</span><span class="p">):</span>
+            <span class="c1"># use mask unit attention until first layer that uses self attention</span>
+            <span class="k">if</span> <span class="n">stage</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;self_attention&quot;</span><span class="p">,</span> <span class="kc">False</span><span class="p">):</span>
+                <span class="k">break</span>
+            <span class="nb">print</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;Stage: </span><span class="si">{</span><span class="n">stage_num</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
+            <span class="k">for</span> <span class="n">block</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">stage</span><span class="p">[</span><span class="s2">&quot;repeat&quot;</span><span class="p">]):</span>
+                <span class="n">is_last</span> <span class="o">=</span> <span class="n">block</span> <span class="o">==</span> <span class="n">stage</span><span class="p">[</span><span class="s2">&quot;repeat&quot;</span><span class="p">]</span> <span class="o">-</span> <span class="mi">1</span>
+                <span class="c1"># do spatial pooling within mask unit on last block of stage</span>
+                <span class="n">q_stride</span> <span class="o">=</span> <span class="n">stage</span><span class="p">[</span><span class="s2">&quot;q_stride&quot;</span><span class="p">]</span> <span class="k">if</span> <span class="n">is_last</span> <span class="k">else</span> <span class="p">[</span><span class="mi">1</span><span class="p">]</span> <span class="o">*</span> <span class="n">spatial_dims</span>
+
+                <span class="c1"># double embedding dimension in last block of stage</span>
+                <span class="n">dim_in</span> <span class="o">=</span> <span class="n">emb_dim</span> <span class="o">*</span> <span class="p">(</span><span class="mi">2</span><span class="o">**</span><span class="n">stage_num</span><span class="p">)</span>
+                <span class="n">dim_out</span> <span class="o">=</span> <span class="n">dim_in</span> <span class="k">if</span> <span class="ow">not</span> <span class="n">is_last</span> <span class="k">else</span> <span class="n">dim_in</span> <span class="o">*</span> <span class="mi">2</span>
+                <span class="nb">print</span><span class="p">(</span>
+                    <span class="sa">f</span><span class="s2">&quot;</span><span class="se">\t</span><span class="s2">Block </span><span class="si">{</span><span class="n">block</span><span class="si">}</span><span class="s2">:</span><span class="se">\t\t</span><span class="s2">dim_in: </span><span class="si">{</span><span class="n">dim_in</span><span class="si">}</span><span class="s2">, dim_out: </span><span class="si">{</span><span class="n">dim_out</span><span class="si">}</span><span class="s2">, num_heads: </span><span class="si">{</span><span class="n">stage</span><span class="p">[</span><span class="s1">&#39;num_heads&#39;</span><span class="p">]</span><span class="si">}</span><span class="s2">, q_stride: </span><span class="si">{</span><span class="n">q_stride</span><span class="si">}</span><span class="s2">, patches_per_mask_unit: </span><span class="si">{</span><span class="n">patches_per_mask_unit</span><span class="si">}</span><span class="s2">&quot;</span>
+                <span class="p">)</span>
+                <span class="n">transformer</span><span class="o">.</span><span class="n">append</span><span class="p">(</span>
+                    <span class="n">HieraBlock</span><span class="p">(</span>
+                        <span class="n">dim</span><span class="o">=</span><span class="n">dim_in</span><span class="p">,</span>
+                        <span class="n">dim_out</span><span class="o">=</span><span class="n">dim_out</span><span class="p">,</span>
+                        <span class="n">heads</span><span class="o">=</span><span class="n">stage</span><span class="p">[</span><span class="s2">&quot;num_heads&quot;</span><span class="p">],</span>
+                        <span class="n">spatial_dims</span><span class="o">=</span><span class="n">spatial_dims</span><span class="p">,</span>
+                        <span class="n">q_stride</span><span class="o">=</span><span class="n">q_stride</span><span class="p">,</span>
+                        <span class="n">patches_per_mask_unit</span><span class="o">=</span><span class="n">patches_per_mask_unit</span><span class="p">,</span>
+                    <span class="p">)</span>
+                <span class="p">)</span>
+                <span class="k">if</span> <span class="n">is_last</span><span class="p">:</span>
+                    <span class="c1"># save the block before the spatial pooling unless it&#39;s the final stage</span>
+                    <span class="n">save_block</span> <span class="o">=</span> <span class="p">(</span>
+                        <span class="n">num_blocks</span> <span class="o">-</span> <span class="mi">1</span> <span class="k">if</span> <span class="n">stage_num</span> <span class="o">&lt;</span> <span class="nb">len</span><span class="p">(</span><span class="n">architecture</span><span class="p">)</span> <span class="o">-</span> <span class="mi">1</span> <span class="k">else</span> <span class="n">num_blocks</span>
+                    <span class="p">)</span>
+                    <span class="bp">self</span><span class="o">.</span><span class="n">save_block_idxs</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">save_block</span><span class="p">)</span>
+                    <span class="bp">self</span><span class="o">.</span><span class="n">save_block_dims</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">dim_in</span><span class="p">)</span>
+
+                    <span class="c1"># create a spatial merger for combining tokens pre-downsampling, last stage doesn&#39;t need merging since it has expected num channels, spatial shape</span>
+                    <span class="bp">self</span><span class="o">.</span><span class="n">spatial_mergers</span><span class="p">[</span><span class="sa">f</span><span class="s2">&quot;block_</span><span class="si">{</span><span class="n">save_block</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="p">(</span>
+                        <span class="n">SpatialMerger</span><span class="p">(</span>
+                            <span class="n">patches_per_mask_unit</span><span class="p">,</span>
+                            <span class="n">dim_in</span><span class="p">,</span>
+                            <span class="bp">self</span><span class="o">.</span><span class="n">final_dim</span><span class="p">,</span>
+                            <span class="n">spatial_dims</span><span class="o">=</span><span class="n">spatial_dims</span><span class="p">,</span>
+                        <span class="p">)</span>
+                        <span class="k">if</span> <span class="n">stage_num</span> <span class="o">&lt;</span> <span class="nb">len</span><span class="p">(</span><span class="n">architecture</span><span class="p">)</span> <span class="o">-</span> <span class="mi">1</span>
+                        <span class="k">else</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Identity</span><span class="p">()</span>
+                    <span class="p">)</span>
+
+                    <span class="c1"># at end of each layer, patches per mask unit is reduced as we pool spatially within mask units</span>
+                    <span class="n">patches_per_mask_unit</span> <span class="o">=</span> <span class="n">patches_per_mask_unit</span> <span class="o">//</span> <span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">(</span><span class="n">stage</span><span class="p">[</span><span class="s2">&quot;q_stride&quot;</span><span class="p">])</span>
+                <span class="n">num_blocks</span> <span class="o">+=</span> <span class="mi">1</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">mask_unit_transformer</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Sequential</span><span class="p">(</span><span class="o">*</span><span class="n">transformer</span><span class="p">)</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">save_block_dims</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">final_dim</span><span class="p">)</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">save_block_dims</span><span class="o">.</span><span class="n">reverse</span><span class="p">()</span>
+
+        <span class="bp">self</span><span class="o">.</span><span class="n">self_attention_transformer</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Sequential</span><span class="p">(</span>
+            <span class="o">*</span><span class="p">[</span><span class="n">Block</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">final_dim</span><span class="p">,</span> <span class="n">stage</span><span class="p">[</span><span class="s2">&quot;num_heads&quot;</span><span class="p">])</span> <span class="k">for</span> <span class="n">_</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">stage</span><span class="p">[</span><span class="s2">&quot;repeat&quot;</span><span class="p">])]</span>
+        <span class="p">)</span>
+
+        <span class="bp">self</span><span class="o">.</span><span class="n">layer_norm</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">LayerNorm</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">final_dim</span><span class="p">)</span>
+
+<div class="viewcode-block" id="HieraEncoder.forward"><a class="viewcode-back" href="../../../../cyto_dl.nn.vits.encoder.html#cyto_dl.nn.vits.encoder.HieraEncoder.forward">[docs]</a>    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">img</span><span class="p">,</span> <span class="n">mask_ratio</span><span class="p">):</span>
+        <span class="n">patches</span><span class="p">,</span> <span class="n">mask</span><span class="p">,</span> <span class="n">forward_indexes</span><span class="p">,</span> <span class="n">backward_indexes</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">patchify</span><span class="p">(</span><span class="n">img</span><span class="p">,</span> <span class="n">mask_ratio</span><span class="p">)</span>
+
+        <span class="c1"># mask unit attention</span>
+        <span class="n">mask_unit_embeddings</span> <span class="o">=</span> <span class="mf">0.0</span>
+        <span class="n">save_layers</span> <span class="o">=</span> <span class="p">[]</span>
+        <span class="k">for</span> <span class="n">i</span><span class="p">,</span> <span class="n">block</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">mask_unit_transformer</span><span class="p">):</span>
+            <span class="n">patches</span> <span class="o">=</span> <span class="n">block</span><span class="p">(</span><span class="n">patches</span><span class="p">)</span>
+            <span class="k">if</span> <span class="n">i</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">save_block_idxs</span><span class="p">:</span>
+                <span class="n">mask_unit_embeddings</span> <span class="o">+=</span> <span class="bp">self</span><span class="o">.</span><span class="n">spatial_mergers</span><span class="p">[</span><span class="sa">f</span><span class="s2">&quot;block_</span><span class="si">{</span><span class="n">i</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">](</span><span class="n">patches</span><span class="p">)</span>
+                <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">save_layers</span><span class="p">:</span>
+                    <span class="n">save_layers</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">patches</span><span class="p">)</span>
+
+        <span class="c1"># combine mask units and tokens for full self attention transformer</span>
+        <span class="n">mask_unit_embeddings</span> <span class="o">=</span> <span class="n">rearrange</span><span class="p">(</span><span class="n">mask_unit_embeddings</span><span class="p">,</span> <span class="s2">&quot;b n_mu t c -&gt; b (n_mu t) c&quot;</span><span class="p">)</span>
+        <span class="n">mask_unit_embeddings</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">self_attention_transformer</span><span class="p">(</span><span class="n">mask_unit_embeddings</span><span class="p">)</span>
+        <span class="n">mask_unit_embeddings</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">layer_norm</span><span class="p">(</span><span class="n">mask_unit_embeddings</span><span class="p">)</span>
+        <span class="n">mask_unit_embeddings</span> <span class="o">=</span> <span class="n">rearrange</span><span class="p">(</span><span class="n">mask_unit_embeddings</span><span class="p">,</span> <span class="s2">&quot;b t c -&gt; t b c&quot;</span><span class="p">)</span>
+
+        <span class="k">return</span> <span class="n">mask_unit_embeddings</span><span class="p">,</span> <span class="n">mask</span><span class="p">,</span> <span class="n">forward_indexes</span><span class="p">,</span> <span class="n">backward_indexes</span>  <span class="c1"># , save_layers</span></div></div>
+</pre></div>
+        </article>
+      </div>
+      <footer>
+        
+        <div class="related-pages">
+          
+          
+        </div>
+        <div class="bottom-of-page">
+          <div class="left-details">
+            <div class="copyright">
+                Copyright &#169; 2023
+            </div>
+            Made with <a href="https://www.sphinx-doc.org/">Sphinx</a> and <a class="muted-link" href="https://pradyunsg.me">@pradyunsg</a>'s
+            
+            <a href="https://github.com/pradyunsg/furo">Furo</a>
+            
+          </div>
+          <div class="right-details">
+            
+          </div>
+        </div>
+        
+      </footer>
+    </div>
+    <aside class="toc-drawer no-toc">
+      
+      
+      
+    </aside>
+  </div>
+</div><script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
+    <script src="../../../../_static/doctools.js"></script>
+    <script src="../../../../_static/sphinx_highlight.js"></script>
+    <script src="../../../../_static/scripts/furo.js"></script>
+    </body>
+</html>
\ No newline at end of file
diff --git a/_modules/cyto_dl/nn/vits/mae.html b/_modules/cyto_dl/nn/vits/mae.html
index 15cacbcb3..cb9af9472 100644
--- a/_modules/cyto_dl/nn/vits/mae.html
+++ b/_modules/cyto_dl/nn/vits/mae.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
@@ -435,221 +442,64 @@
           <h1>Source code for cyto_dl.nn.vits.mae</h1><div class="highlight"><pre>
 <span></span><span class="c1"># modified from https://github.com/IcarusWizard/MAE/blob/main/model.py#L124</span>
 
-<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">List</span><span class="p">,</span> <span class="n">Optional</span>
+<span class="kn">from</span> <span class="nn">abc</span> <span class="kn">import</span> <span class="n">ABC</span><span class="p">,</span> <span class="n">abstractmethod</span>
+<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Dict</span><span class="p">,</span> <span class="n">List</span><span class="p">,</span> <span class="n">Optional</span><span class="p">,</span> <span class="n">Union</span>
 
 <span class="kn">import</span> <span class="nn">numpy</span> <span class="k">as</span> <span class="nn">np</span>
 <span class="kn">import</span> <span class="nn">torch</span>
-<span class="kn">import</span> <span class="nn">torch.nn</span> <span class="k">as</span> <span class="nn">nn</span>
-<span class="kn">from</span> <span class="nn">einops</span> <span class="kn">import</span> <span class="n">rearrange</span>
-<span class="kn">from</span> <span class="nn">einops.layers.torch</span> <span class="kn">import</span> <span class="n">Rearrange</span>
-<span class="kn">from</span> <span class="nn">timm.models.layers</span> <span class="kn">import</span> <span class="n">trunc_normal_</span>
-<span class="kn">from</span> <span class="nn">timm.models.vision_transformer</span> <span class="kn">import</span> <span class="n">Block</span>
 
-<span class="kn">from</span> <span class="nn">cyto_dl.nn.vits.blocks</span> <span class="kn">import</span> <span class="n">IntermediateWeigher</span><span class="p">,</span> <span class="n">Patchify</span>
-<span class="kn">from</span> <span class="nn">cyto_dl.nn.vits.cross_mae</span> <span class="kn">import</span> <span class="n">CrossMAE_Decoder</span>
-<span class="kn">from</span> <span class="nn">cyto_dl.nn.vits.utils</span> <span class="kn">import</span> <span class="n">get_positional_embedding</span><span class="p">,</span> <span class="n">take_indexes</span>
+<span class="kn">from</span> <span class="nn">cyto_dl.nn.vits.decoder</span> <span class="kn">import</span> <span class="n">CrossMAE_Decoder</span><span class="p">,</span> <span class="n">MAE_Decoder</span>
+<span class="kn">from</span> <span class="nn">cyto_dl.nn.vits.encoder</span> <span class="kn">import</span> <span class="n">HieraEncoder</span><span class="p">,</span> <span class="n">MAE_Encoder</span>
+<span class="kn">from</span> <span class="nn">cyto_dl.nn.vits.utils</span> <span class="kn">import</span> <span class="n">match_tuple_dimensions</span>
 
 
-<div class="viewcode-block" id="MAE_Encoder"><a class="viewcode-back" href="../../../../cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.MAE_Encoder">[docs]</a><span class="k">class</span> <span class="nc">MAE_Encoder</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
+<div class="viewcode-block" id="MAE_Base"><a class="viewcode-back" href="../../../../cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.MAE_Base">[docs]</a><span class="k">class</span> <span class="nc">MAE_Base</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">,</span> <span class="n">ABC</span><span class="p">):</span>
     <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span>
-        <span class="bp">self</span><span class="p">,</span>
-        <span class="n">num_patches</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span>
-        <span class="n">spatial_dims</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">3</span><span class="p">,</span>
-        <span class="n">base_patch_size</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="p">(</span><span class="mi">16</span><span class="p">,</span> <span class="mi">16</span><span class="p">,</span> <span class="mi">16</span><span class="p">),</span>
-        <span class="n">emb_dim</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">192</span><span class="p">,</span>
-        <span class="n">num_layer</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">12</span><span class="p">,</span>
-        <span class="n">num_head</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">3</span><span class="p">,</span>
-        <span class="n">context_pixels</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">]]</span> <span class="o">=</span> <span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">],</span>
-        <span class="n">input_channels</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">1</span><span class="p">,</span>
-        <span class="n">n_intermediate_weights</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="o">-</span><span class="mi">1</span><span class="p">,</span>
-    <span class="p">)</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
-<span class="w">        </span><span class="sd">&quot;&quot;&quot;</span>
-<span class="sd">        Parameters</span>
-<span class="sd">        ----------</span>
-<span class="sd">        num_patches: List[int]</span>
-<span class="sd">            Number of patches in each dimension</span>
-<span class="sd">        spatial_dims: int</span>
-<span class="sd">            Number of spatial dimensions</span>
-<span class="sd">        base_patch_size: List[int]</span>
-<span class="sd">            Size of each patch</span>
-<span class="sd">        emb_dim: int</span>
-<span class="sd">            Dimension of embedding</span>
-<span class="sd">        num_layer: int</span>
-<span class="sd">            Number of transformer layers</span>
-<span class="sd">        num_head: int</span>
-<span class="sd">            Number of heads in transformer</span>
-<span class="sd">        context_pixels: List[int]</span>
-<span class="sd">            Number of extra pixels around each patch to include in convolutional embedding to encoder dimension.</span>
-<span class="sd">        input_channels: int</span>
-<span class="sd">            Number of input channels</span>
-<span class="sd">        weight_intermediates: bool</span>
-<span class="sd">            Whether to output linear combination of intermediate layers as final output like CrossMAE</span>
-<span class="sd">        &quot;&quot;&quot;</span>
+        <span class="bp">self</span><span class="p">,</span> <span class="n">spatial_dims</span><span class="p">,</span> <span class="n">num_patches</span><span class="p">,</span> <span class="n">patch_size</span><span class="p">,</span> <span class="n">mask_ratio</span><span class="p">,</span> <span class="n">features_only</span><span class="p">,</span> <span class="n">context_pixels</span>
+    <span class="p">):</span>
         <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>
-        <span class="bp">self</span><span class="o">.</span><span class="n">cls_token</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Parameter</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="n">emb_dim</span><span class="p">))</span>
-        <span class="bp">self</span><span class="o">.</span><span class="n">patchify</span> <span class="o">=</span> <span class="n">Patchify</span><span class="p">(</span>
-            <span class="n">base_patch_size</span><span class="p">,</span> <span class="n">emb_dim</span><span class="p">,</span> <span class="n">num_patches</span><span class="p">,</span> <span class="n">spatial_dims</span><span class="p">,</span> <span class="n">context_pixels</span><span class="p">,</span> <span class="n">input_channels</span>
-        <span class="p">)</span>
-        <span class="n">weight_intermediates</span> <span class="o">=</span> <span class="n">n_intermediate_weights</span> <span class="o">&gt;</span> <span class="mi">0</span>
-        <span class="k">if</span> <span class="n">weight_intermediates</span><span class="p">:</span>
-            <span class="bp">self</span><span class="o">.</span><span class="n">transformer</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">ModuleList</span><span class="p">(</span>
-                <span class="p">[</span><span class="n">Block</span><span class="p">(</span><span class="n">emb_dim</span><span class="p">,</span> <span class="n">num_head</span><span class="p">)</span> <span class="k">for</span> <span class="n">_</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">num_layer</span><span class="p">)]</span>
-            <span class="p">)</span>
-        <span class="k">else</span><span class="p">:</span>
-            <span class="bp">self</span><span class="o">.</span><span class="n">transformer</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Sequential</span><span class="p">(</span>
-                <span class="o">*</span><span class="p">[</span><span class="n">Block</span><span class="p">(</span><span class="n">emb_dim</span><span class="p">,</span> <span class="n">num_head</span><span class="p">)</span> <span class="k">for</span> <span class="n">_</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">num_layer</span><span class="p">)]</span>
-            <span class="p">)</span>
-
-        <span class="bp">self</span><span class="o">.</span><span class="n">layer_norm</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">LayerNorm</span><span class="p">(</span><span class="n">emb_dim</span><span class="p">)</span>
-
-        <span class="bp">self</span><span class="o">.</span><span class="n">intermediate_weighter</span> <span class="o">=</span> <span class="p">(</span>
-            <span class="n">IntermediateWeigher</span><span class="p">(</span><span class="n">num_layer</span><span class="p">,</span> <span class="n">emb_dim</span><span class="p">,</span> <span class="n">n_intermediate_weights</span><span class="p">)</span>
-            <span class="k">if</span> <span class="n">weight_intermediates</span>
-            <span class="k">else</span> <span class="kc">None</span>
+        <span class="n">num_patches</span><span class="p">,</span> <span class="n">patch_size</span><span class="p">,</span> <span class="n">context_pixels</span> <span class="o">=</span> <span class="n">match_tuple_dimensions</span><span class="p">(</span>
+            <span class="n">spatial_dims</span><span class="p">,</span> <span class="p">[</span><span class="n">num_patches</span><span class="p">,</span> <span class="n">patch_size</span><span class="p">,</span> <span class="n">context_pixels</span><span class="p">]</span>
         <span class="p">)</span>
-        <span class="bp">self</span><span class="o">.</span><span class="n">init_weight</span><span class="p">()</span>
-
-<div class="viewcode-block" id="MAE_Encoder.init_weight"><a class="viewcode-back" href="../../../../cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.MAE_Encoder.init_weight">[docs]</a>    <span class="k">def</span> <span class="nf">init_weight</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
-        <span class="n">trunc_normal_</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">cls_token</span><span class="p">,</span> <span class="n">std</span><span class="o">=</span><span class="mf">0.02</span><span class="p">)</span></div>
-
-<div class="viewcode-block" id="MAE_Encoder.forward"><a class="viewcode-back" href="../../../../cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.MAE_Encoder.forward">[docs]</a>    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">img</span><span class="p">,</span> <span class="n">mask_ratio</span><span class="o">=</span><span class="mf">0.75</span><span class="p">):</span>
-        <span class="n">patches</span><span class="p">,</span> <span class="n">mask</span><span class="p">,</span> <span class="n">forward_indexes</span><span class="p">,</span> <span class="n">backward_indexes</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">patchify</span><span class="p">(</span><span class="n">img</span><span class="p">,</span> <span class="n">mask_ratio</span><span class="p">)</span>
-        <span class="n">patches</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">cat</span><span class="p">([</span><span class="bp">self</span><span class="o">.</span><span class="n">cls_token</span><span class="o">.</span><span class="n">expand</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="n">patches</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class="o">-</span><span class="mi">1</span><span class="p">),</span> <span class="n">patches</span><span class="p">],</span> <span class="n">dim</span><span class="o">=</span><span class="mi">0</span><span class="p">)</span>
-        <span class="n">patches</span> <span class="o">=</span> <span class="n">rearrange</span><span class="p">(</span><span class="n">patches</span><span class="p">,</span> <span class="s2">&quot;t b c -&gt; b t c&quot;</span><span class="p">)</span>
-
-        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">intermediate_weighter</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
-            <span class="n">intermediates</span> <span class="o">=</span> <span class="p">[</span><span class="n">patches</span><span class="p">]</span>
-            <span class="k">for</span> <span class="n">block</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">transformer</span><span class="p">:</span>
-                <span class="n">patches</span> <span class="o">=</span> <span class="n">block</span><span class="p">(</span><span class="n">patches</span><span class="p">)</span>
-                <span class="n">intermediates</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">patches</span><span class="p">)</span>
-            <span class="n">features</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">layer_norm</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">intermediate_weighter</span><span class="p">(</span><span class="n">intermediates</span><span class="p">))</span>
-            <span class="n">features</span> <span class="o">=</span> <span class="n">rearrange</span><span class="p">(</span><span class="n">features</span><span class="p">,</span> <span class="s2">&quot;n b t c -&gt; n t b c&quot;</span><span class="p">)</span>
-        <span class="k">else</span><span class="p">:</span>
-            <span class="n">features</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">layer_norm</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">transformer</span><span class="p">(</span><span class="n">patches</span><span class="p">))</span>
-            <span class="n">features</span> <span class="o">=</span> <span class="n">rearrange</span><span class="p">(</span><span class="n">features</span><span class="p">,</span> <span class="s2">&quot;b t c -&gt; t b c&quot;</span><span class="p">)</span>
-        <span class="k">if</span> <span class="n">mask_ratio</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>
-            <span class="k">return</span> <span class="n">features</span><span class="p">,</span> <span class="n">mask</span><span class="p">,</span> <span class="n">forward_indexes</span><span class="p">,</span> <span class="n">backward_indexes</span>
-        <span class="k">return</span> <span class="n">features</span></div></div>
-
-
-<div class="viewcode-block" id="MAE_Decoder"><a class="viewcode-back" href="../../../../cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.MAE_Decoder">[docs]</a><span class="k">class</span> <span class="nc">MAE_Decoder</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
-    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span>
-        <span class="bp">self</span><span class="p">,</span>
-        <span class="n">num_patches</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span>
-        <span class="n">spatial_dims</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">3</span><span class="p">,</span>
-        <span class="n">base_patch_size</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">]]</span> <span class="o">=</span> <span class="p">[</span><span class="mi">4</span><span class="p">,</span> <span class="mi">8</span><span class="p">,</span> <span class="mi">8</span><span class="p">],</span>
-        <span class="n">enc_dim</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">768</span><span class="p">,</span>
-        <span class="n">emb_dim</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">192</span><span class="p">,</span>
-        <span class="n">num_layer</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">4</span><span class="p">,</span>
-        <span class="n">num_head</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">3</span><span class="p">,</span>
-        <span class="n">learnable_pos_embedding</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">True</span><span class="p">,</span>
-    <span class="p">)</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
-<span class="w">        </span><span class="sd">&quot;&quot;&quot;</span>
-<span class="sd">        Parameters</span>
-<span class="sd">        ----------</span>
-<span class="sd">        num_patches: List[int]</span>
-<span class="sd">            Number of patches in each dimension</span>
-<span class="sd">        base_patch_size: Tuple[int]</span>
-<span class="sd">            Size of each patch</span>
-<span class="sd">        enc_dim: int</span>
-<span class="sd">            Dimension of encoder</span>
-<span class="sd">        emb_dim: int</span>
-<span class="sd">            Dimension of decoder</span>
-<span class="sd">        num_layer: int</span>
-<span class="sd">            Number of transformer layers</span>
-<span class="sd">        num_head: int</span>
-<span class="sd">            Number of heads in transformer</span>
-<span class="sd">        learnable_pos_embedding: bool</span>
-<span class="sd">            If True, learnable positional embeddings are used. If False, fixed sin/cos positional embeddings. Empirically, fixed positional embeddings work better for brightfield images.</span>
-<span class="sd">        &quot;&quot;&quot;</span>
-        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>
-        <span class="bp">self</span><span class="o">.</span><span class="n">projection_norm</span> <span class="o">=</span> <span class="n">nn</span><span class="o">.</span><span class="n">LayerNorm</span><span class="p">(</span><span class="n">emb_dim</span><span class="p">)</span>
-        <span class="bp">self</span><span class="o">.</span><span class="n">projection</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Linear</span><span class="p">(</span><span class="n">enc_dim</span><span class="p">,</span> <span class="n">emb_dim</span><span class="p">)</span>
-        <span class="bp">self</span><span class="o">.</span><span class="n">mask_token</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Parameter</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="n">emb_dim</span><span class="p">))</span>
 
-        <span class="bp">self</span><span class="o">.</span><span class="n">pos_embedding</span> <span class="o">=</span> <span class="n">get_positional_embedding</span><span class="p">(</span>
-            <span class="n">num_patches</span><span class="p">,</span> <span class="n">emb_dim</span><span class="p">,</span> <span class="n">learnable</span><span class="o">=</span><span class="n">learnable_pos_embedding</span>
-        <span class="p">)</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">spatial_dims</span> <span class="o">=</span> <span class="n">spatial_dims</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">num_patches</span> <span class="o">=</span> <span class="n">num_patches</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">patch_size</span> <span class="o">=</span> <span class="n">patch_size</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">mask_ratio</span> <span class="o">=</span> <span class="n">mask_ratio</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">features_only</span> <span class="o">=</span> <span class="n">features_only</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">context_pixels</span> <span class="o">=</span> <span class="n">context_pixels</span>
 
-        <span class="bp">self</span><span class="o">.</span><span class="n">transformer</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Sequential</span><span class="p">(</span>
-            <span class="o">*</span><span class="p">[</span><span class="n">Block</span><span class="p">(</span><span class="n">emb_dim</span><span class="p">,</span> <span class="n">num_head</span><span class="p">)</span> <span class="k">for</span> <span class="n">_</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">num_layer</span><span class="p">)]</span>
-        <span class="p">)</span>
-        <span class="n">out_dim</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">prod</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">as_tensor</span><span class="p">(</span><span class="n">base_patch_size</span><span class="p">))</span><span class="o">.</span><span class="n">item</span><span class="p">()</span>
-        <span class="bp">self</span><span class="o">.</span><span class="n">head_norm</span> <span class="o">=</span> <span class="n">nn</span><span class="o">.</span><span class="n">LayerNorm</span><span class="p">(</span><span class="n">out_dim</span><span class="p">)</span>
-        <span class="bp">self</span><span class="o">.</span><span class="n">head</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Linear</span><span class="p">(</span><span class="n">emb_dim</span><span class="p">,</span> <span class="n">out_dim</span><span class="p">)</span>
-        <span class="bp">self</span><span class="o">.</span><span class="n">num_patches</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">as_tensor</span><span class="p">(</span><span class="n">num_patches</span><span class="p">)</span>
-
-        <span class="k">if</span> <span class="n">spatial_dims</span> <span class="o">==</span> <span class="mi">3</span><span class="p">:</span>
-            <span class="bp">self</span><span class="o">.</span><span class="n">patch2img</span> <span class="o">=</span> <span class="n">Rearrange</span><span class="p">(</span>
-                <span class="s2">&quot;(n_patch_z n_patch_y n_patch_x) b (c patch_size_z patch_size_y patch_size_x) -&gt;  b c (n_patch_z patch_size_z) (n_patch_y patch_size_y) (n_patch_x patch_size_x)&quot;</span><span class="p">,</span>
-                <span class="n">n_patch_z</span><span class="o">=</span><span class="n">num_patches</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span>
-                <span class="n">n_patch_y</span><span class="o">=</span><span class="n">num_patches</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span>
-                <span class="n">n_patch_x</span><span class="o">=</span><span class="n">num_patches</span><span class="p">[</span><span class="mi">2</span><span class="p">],</span>
-                <span class="n">patch_size_z</span><span class="o">=</span><span class="n">base_patch_size</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span>
-                <span class="n">patch_size_y</span><span class="o">=</span><span class="n">base_patch_size</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span>
-                <span class="n">patch_size_x</span><span class="o">=</span><span class="n">base_patch_size</span><span class="p">[</span><span class="mi">2</span><span class="p">],</span>
-            <span class="p">)</span>
-        <span class="k">elif</span> <span class="n">spatial_dims</span> <span class="o">==</span> <span class="mi">2</span><span class="p">:</span>
-            <span class="bp">self</span><span class="o">.</span><span class="n">patch2img</span> <span class="o">=</span> <span class="n">Rearrange</span><span class="p">(</span>
-                <span class="s2">&quot;(n_patch_y n_patch_x) b (c patch_size_y patch_size_x) -&gt;  b c (n_patch_y patch_size_y) (n_patch_x patch_size_x)&quot;</span><span class="p">,</span>
-                <span class="n">n_patch_y</span><span class="o">=</span><span class="n">num_patches</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span>
-                <span class="n">n_patch_x</span><span class="o">=</span><span class="n">num_patches</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span>
-                <span class="n">patch_size_y</span><span class="o">=</span><span class="n">base_patch_size</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span>
-                <span class="n">patch_size_x</span><span class="o">=</span><span class="n">base_patch_size</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span>
-            <span class="p">)</span>
-
-        <span class="bp">self</span><span class="o">.</span><span class="n">init_weight</span><span class="p">()</span>
-
-<div class="viewcode-block" id="MAE_Decoder.init_weight"><a class="viewcode-back" href="../../../../cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.MAE_Decoder.init_weight">[docs]</a>    <span class="k">def</span> <span class="nf">init_weight</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
-        <span class="n">trunc_normal_</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">mask_token</span><span class="p">,</span> <span class="n">std</span><span class="o">=</span><span class="mf">0.02</span><span class="p">)</span></div>
-
-<div class="viewcode-block" id="MAE_Decoder.forward"><a class="viewcode-back" href="../../../../cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.MAE_Decoder.forward">[docs]</a>    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">features</span><span class="p">,</span> <span class="n">forward_indexes</span><span class="p">,</span> <span class="n">backward_indexes</span><span class="p">):</span>
-        <span class="c1"># project from encoder dimension to decoder dimension</span>
-        <span class="n">features</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">projection_norm</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">projection</span><span class="p">(</span><span class="n">features</span><span class="p">))</span>
-
-        <span class="n">backward_indexes</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">cat</span><span class="p">(</span>
-            <span class="p">[</span>
-                <span class="n">torch</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span>
-                    <span class="mi">1</span><span class="p">,</span> <span class="n">backward_indexes</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class="n">device</span><span class="o">=</span><span class="n">backward_indexes</span><span class="o">.</span><span class="n">device</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">torch</span><span class="o">.</span><span class="n">long</span>
-                <span class="p">),</span>
-                <span class="n">backward_indexes</span> <span class="o">+</span> <span class="mi">1</span><span class="p">,</span>
-            <span class="p">],</span>
-            <span class="n">dim</span><span class="o">=</span><span class="mi">0</span><span class="p">,</span>
-        <span class="p">)</span>
-        <span class="c1"># fill in masked regions</span>
-        <span class="n">features</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">cat</span><span class="p">(</span>
-            <span class="p">[</span>
-                <span class="n">features</span><span class="p">,</span>
-                <span class="bp">self</span><span class="o">.</span><span class="n">mask_token</span><span class="o">.</span><span class="n">expand</span><span class="p">(</span>
-                    <span class="n">backward_indexes</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">-</span> <span class="n">features</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">features</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class="o">-</span><span class="mi">1</span>
-                <span class="p">),</span>
-            <span class="p">],</span>
-            <span class="n">dim</span><span class="o">=</span><span class="mi">0</span><span class="p">,</span>
-        <span class="p">)</span>
-        <span class="c1"># unshuffle to original positions</span>
-        <span class="n">features</span> <span class="o">=</span> <span class="n">take_indexes</span><span class="p">(</span><span class="n">features</span><span class="p">,</span> <span class="n">backward_indexes</span><span class="p">)</span>
-        <span class="n">features</span> <span class="o">=</span> <span class="n">features</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">pos_embedding</span>
+    <span class="c1"># encoder and decoder must be defined in subclasses</span>
+    <span class="nd">@property</span>
+    <span class="nd">@abstractmethod</span>
+    <span class="k">def</span> <span class="nf">encoder</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
+        <span class="k">pass</span>
 
-        <span class="c1"># decode</span>
-        <span class="n">features</span> <span class="o">=</span> <span class="n">rearrange</span><span class="p">(</span><span class="n">features</span><span class="p">,</span> <span class="s2">&quot;t b c -&gt; b t c&quot;</span><span class="p">)</span>
-        <span class="n">features</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">transformer</span><span class="p">(</span><span class="n">features</span><span class="p">)</span>
-        <span class="n">features</span> <span class="o">=</span> <span class="n">rearrange</span><span class="p">(</span><span class="n">features</span><span class="p">,</span> <span class="s2">&quot;b t c -&gt; t b c&quot;</span><span class="p">)</span>
-        <span class="n">features</span> <span class="o">=</span> <span class="n">features</span><span class="p">[</span><span class="mi">1</span><span class="p">:]</span>  <span class="c1"># remove global feature</span>
+    <span class="nd">@property</span>
+    <span class="nd">@abstractmethod</span>
+    <span class="k">def</span> <span class="nf">decoder</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
+        <span class="k">pass</span>
 
-        <span class="c1"># (npatches x npatches x npatches) b (emb dim) -&gt; (npatches* npatches * npatches) b (z y x)</span>
-        <span class="n">patches</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">head_norm</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">head</span><span class="p">(</span><span class="n">features</span><span class="p">))</span>
+<div class="viewcode-block" id="MAE_Base.init_encoder"><a class="viewcode-back" href="../../../../cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.MAE_Base.init_encoder">[docs]</a>    <span class="k">def</span> <span class="nf">init_encoder</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
+        <span class="k">raise</span> <span class="ne">NotImplementedError</span></div>
 
-        <span class="c1"># patches to image</span>
-        <span class="n">img</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">patch2img</span><span class="p">(</span><span class="n">patches</span><span class="p">)</span>
-        <span class="k">return</span> <span class="n">img</span></div></div>
+<div class="viewcode-block" id="MAE_Base.init_decoder"><a class="viewcode-back" href="../../../../cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.MAE_Base.init_decoder">[docs]</a>    <span class="k">def</span> <span class="nf">init_decoder</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
+        <span class="k">raise</span> <span class="ne">NotImplementedError</span></div>
+
+<div class="viewcode-block" id="MAE_Base.forward"><a class="viewcode-back" href="../../../../cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.MAE_Base.forward">[docs]</a>    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">img</span><span class="p">):</span>
+        <span class="n">features</span><span class="p">,</span> <span class="n">mask</span><span class="p">,</span> <span class="n">forward_indexes</span><span class="p">,</span> <span class="n">backward_indexes</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">encoder</span><span class="p">(</span><span class="n">img</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">mask_ratio</span><span class="p">)</span>
+        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">features_only</span><span class="p">:</span>
+            <span class="k">return</span> <span class="n">features</span>
+        <span class="n">predicted_img</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">decoder</span><span class="p">(</span><span class="n">features</span><span class="p">,</span> <span class="n">forward_indexes</span><span class="p">,</span> <span class="n">backward_indexes</span><span class="p">)</span>
+        <span class="k">return</span> <span class="n">predicted_img</span><span class="p">,</span> <span class="n">mask</span></div></div>
 
 
-<div class="viewcode-block" id="MAE_ViT"><a class="viewcode-back" href="../../../../cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.MAE_ViT">[docs]</a><span class="k">class</span> <span class="nc">MAE_ViT</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
+<div class="viewcode-block" id="MAE"><a class="viewcode-back" href="../../../../cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.MAE">[docs]</a><span class="k">class</span> <span class="nc">MAE</span><span class="p">(</span><span class="n">MAE_Base</span><span class="p">):</span>
     <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span>
         <span class="bp">self</span><span class="p">,</span>
         <span class="n">spatial_dims</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">3</span><span class="p">,</span>
-        <span class="n">num_patches</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">]]</span> <span class="o">=</span> <span class="p">[</span><span class="mi">2</span><span class="p">,</span> <span class="mi">32</span><span class="p">,</span> <span class="mi">32</span><span class="p">],</span>
-        <span class="n">base_patch_size</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">]]</span> <span class="o">=</span> <span class="p">[</span><span class="mi">16</span><span class="p">,</span> <span class="mi">16</span><span class="p">,</span> <span class="mi">16</span><span class="p">],</span>
+        <span class="n">num_patches</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">]]</span> <span class="o">=</span> <span class="mi">16</span><span class="p">,</span>
+        <span class="n">patch_size</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">]]</span> <span class="o">=</span> <span class="mi">4</span><span class="p">,</span>
         <span class="n">emb_dim</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">768</span><span class="p">,</span>
         <span class="n">encoder_layer</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">12</span><span class="p">,</span>
         <span class="n">encoder_head</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">8</span><span class="p">,</span>
@@ -658,7 +508,7 @@ <h1>Source code for cyto_dl.nn.vits.mae</h1><div class="highlight"><pre>
         <span class="n">decoder_dim</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">192</span><span class="p">,</span>
         <span class="n">mask_ratio</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mf">0.75</span><span class="p">,</span>
         <span class="n">use_crossmae</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span>
-        <span class="n">context_pixels</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">]]</span> <span class="o">=</span> <span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">],</span>
+        <span class="n">context_pixels</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">]]</span> <span class="o">=</span> <span class="mi">0</span><span class="p">,</span>
         <span class="n">input_channels</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">1</span><span class="p">,</span>
         <span class="n">features_only</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span>
         <span class="n">learnable_pos_embedding</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">True</span><span class="p">,</span>
@@ -670,7 +520,7 @@ <h1>Source code for cyto_dl.nn.vits.mae</h1><div class="highlight"><pre>
 <span class="sd">            Number of spatial dimensions</span>
 <span class="sd">        num_patches: List[int]</span>
 <span class="sd">            Number of patches in each dimension (ZYX order)</span>
-<span class="sd">        base_patch_size: List[int]</span>
+<span class="sd">        patch_size: List[int]</span>
 <span class="sd">            Size of each patch (ZYX order)</span>
 <span class="sd">        emb_dim: int</span>
 <span class="sd">            Dimension of encoder embedding</span>
@@ -695,30 +545,23 @@ <h1>Source code for cyto_dl.nn.vits.mae</h1><div class="highlight"><pre>
 <span class="sd">        learnable_pos_embedding: bool</span>
 <span class="sd">            If True, learnable positional embeddings are used. If False, fixed sin/cos positional embeddings. Empirically, fixed positional embeddings work better for brightfield images.</span>
 <span class="sd">        &quot;&quot;&quot;</span>
-        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>
-        <span class="k">assert</span> <span class="n">spatial_dims</span> <span class="ow">in</span> <span class="p">(</span><span class="mi">2</span><span class="p">,</span> <span class="mi">3</span><span class="p">),</span> <span class="s2">&quot;Spatial dims must be 2 or 3&quot;</span>
-
-        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">num_patches</span><span class="p">,</span> <span class="nb">int</span><span class="p">):</span>
-            <span class="n">num_patches</span> <span class="o">=</span> <span class="p">[</span><span class="n">num_patches</span><span class="p">]</span> <span class="o">*</span> <span class="n">spatial_dims</span>
-        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">base_patch_size</span><span class="p">,</span> <span class="nb">int</span><span class="p">):</span>
-            <span class="n">base_patch_size</span> <span class="o">=</span> <span class="p">[</span><span class="n">base_patch_size</span><span class="p">]</span> <span class="o">*</span> <span class="n">spatial_dims</span>
-
-        <span class="k">assert</span> <span class="nb">len</span><span class="p">(</span><span class="n">num_patches</span><span class="p">)</span> <span class="o">==</span> <span class="n">spatial_dims</span><span class="p">,</span> <span class="s2">&quot;num_patches must be of length spatial_dims&quot;</span>
-        <span class="k">assert</span> <span class="p">(</span>
-            <span class="nb">len</span><span class="p">(</span><span class="n">base_patch_size</span><span class="p">)</span> <span class="o">==</span> <span class="n">spatial_dims</span>
-        <span class="p">),</span> <span class="s2">&quot;base_patch_size must be of length spatial_dims&quot;</span>
-
-        <span class="bp">self</span><span class="o">.</span><span class="n">mask_ratio</span> <span class="o">=</span> <span class="n">mask_ratio</span>
-        <span class="bp">self</span><span class="o">.</span><span class="n">features_only</span> <span class="o">=</span> <span class="n">features_only</span>
+        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span>
+            <span class="n">spatial_dims</span><span class="o">=</span><span class="n">spatial_dims</span><span class="p">,</span>
+            <span class="n">num_patches</span><span class="o">=</span><span class="n">num_patches</span><span class="p">,</span>
+            <span class="n">patch_size</span><span class="o">=</span><span class="n">patch_size</span><span class="p">,</span>
+            <span class="n">mask_ratio</span><span class="o">=</span><span class="n">mask_ratio</span><span class="p">,</span>
+            <span class="n">features_only</span><span class="o">=</span><span class="n">features_only</span><span class="p">,</span>
+            <span class="n">context_pixels</span><span class="o">=</span><span class="n">context_pixels</span><span class="p">,</span>
+        <span class="p">)</span>
 
-        <span class="bp">self</span><span class="o">.</span><span class="n">encoder</span> <span class="o">=</span> <span class="n">MAE_Encoder</span><span class="p">(</span>
-            <span class="n">num_patches</span><span class="p">,</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">_encoder</span> <span class="o">=</span> <span class="n">MAE_Encoder</span><span class="p">(</span>
+            <span class="bp">self</span><span class="o">.</span><span class="n">num_patches</span><span class="p">,</span>
             <span class="n">spatial_dims</span><span class="p">,</span>
-            <span class="n">base_patch_size</span><span class="p">,</span>
+            <span class="bp">self</span><span class="o">.</span><span class="n">patch_size</span><span class="p">,</span>
             <span class="n">emb_dim</span><span class="p">,</span>
             <span class="n">encoder_layer</span><span class="p">,</span>
             <span class="n">encoder_head</span><span class="p">,</span>
-            <span class="n">context_pixels</span><span class="p">,</span>
+            <span class="bp">self</span><span class="o">.</span><span class="n">context_pixels</span><span class="p">,</span>
             <span class="n">input_channels</span><span class="p">,</span>
             <span class="n">n_intermediate_weights</span><span class="o">=-</span><span class="mi">1</span> <span class="k">if</span> <span class="ow">not</span> <span class="n">use_crossmae</span> <span class="k">else</span> <span class="n">decoder_layer</span><span class="p">,</span>
         <span class="p">)</span>
@@ -726,10 +569,10 @@ <h1>Source code for cyto_dl.nn.vits.mae</h1><div class="highlight"><pre>
         <span class="n">decoder_class</span> <span class="o">=</span> <span class="n">MAE_Decoder</span>
         <span class="k">if</span> <span class="n">use_crossmae</span><span class="p">:</span>
             <span class="n">decoder_class</span> <span class="o">=</span> <span class="n">CrossMAE_Decoder</span>
-        <span class="bp">self</span><span class="o">.</span><span class="n">decoder</span> <span class="o">=</span> <span class="n">decoder_class</span><span class="p">(</span>
-            <span class="n">num_patches</span><span class="o">=</span><span class="n">num_patches</span><span class="p">,</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">_decoder</span> <span class="o">=</span> <span class="n">decoder_class</span><span class="p">(</span>
+            <span class="n">num_patches</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">num_patches</span><span class="p">,</span>
             <span class="n">spatial_dims</span><span class="o">=</span><span class="n">spatial_dims</span><span class="p">,</span>
-            <span class="n">base_patch_size</span><span class="o">=</span><span class="n">base_patch_size</span><span class="p">,</span>
+            <span class="n">patch_size</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">patch_size</span><span class="p">,</span>
             <span class="n">enc_dim</span><span class="o">=</span><span class="n">emb_dim</span><span class="p">,</span>
             <span class="n">emb_dim</span><span class="o">=</span><span class="n">decoder_dim</span><span class="p">,</span>
             <span class="n">num_layer</span><span class="o">=</span><span class="n">decoder_layer</span><span class="p">,</span>
@@ -737,12 +580,120 @@ <h1>Source code for cyto_dl.nn.vits.mae</h1><div class="highlight"><pre>
             <span class="n">learnable_pos_embedding</span><span class="o">=</span><span class="n">learnable_pos_embedding</span><span class="p">,</span>
         <span class="p">)</span>
 
-<div class="viewcode-block" id="MAE_ViT.forward"><a class="viewcode-back" href="../../../../cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.MAE_ViT.forward">[docs]</a>    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">img</span><span class="p">):</span>
-        <span class="n">features</span><span class="p">,</span> <span class="n">mask</span><span class="p">,</span> <span class="n">forward_indexes</span><span class="p">,</span> <span class="n">backward_indexes</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">encoder</span><span class="p">(</span><span class="n">img</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">mask_ratio</span><span class="p">)</span>
-        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">features_only</span><span class="p">:</span>
-            <span class="k">return</span> <span class="n">features</span>
-        <span class="n">predicted_img</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">decoder</span><span class="p">(</span><span class="n">features</span><span class="p">,</span> <span class="n">forward_indexes</span><span class="p">,</span> <span class="n">backward_indexes</span><span class="p">)</span>
-        <span class="k">return</span> <span class="n">predicted_img</span><span class="p">,</span> <span class="n">mask</span></div></div>
+    <span class="nd">@property</span>
+    <span class="k">def</span> <span class="nf">encoder</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
+        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_encoder</span>
+
+    <span class="nd">@property</span>
+    <span class="k">def</span> <span class="nf">decoder</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
+        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_decoder</span></div>
+
+
+<div class="viewcode-block" id="HieraMAE"><a class="viewcode-back" href="../../../../cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.HieraMAE">[docs]</a><span class="k">class</span> <span class="nc">HieraMAE</span><span class="p">(</span><span class="n">MAE_Base</span><span class="p">):</span>
+    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span>
+        <span class="bp">self</span><span class="p">,</span>
+        <span class="n">architecture</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="n">Dict</span><span class="p">],</span>
+        <span class="n">spatial_dims</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">3</span><span class="p">,</span>
+        <span class="n">num_patches</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">]]]</span> <span class="o">=</span> <span class="mi">16</span><span class="p">,</span>
+        <span class="n">num_mask_units</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">]]]</span> <span class="o">=</span> <span class="mi">8</span><span class="p">,</span>
+        <span class="n">patch_size</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">]]]</span> <span class="o">=</span> <span class="mi">4</span><span class="p">,</span>
+        <span class="n">emb_dim</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">64</span><span class="p">,</span>
+        <span class="n">decoder_layer</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">4</span><span class="p">,</span>
+        <span class="n">decoder_head</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">8</span><span class="p">,</span>
+        <span class="n">decoder_dim</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">192</span><span class="p">,</span>
+        <span class="n">mask_ratio</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mf">0.75</span><span class="p">,</span>
+        <span class="n">use_crossmae</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span>
+        <span class="n">context_pixels</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">]]</span> <span class="o">=</span> <span class="mi">0</span><span class="p">,</span>
+        <span class="n">input_channels</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">1</span><span class="p">,</span>
+        <span class="n">features_only</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span>
+    <span class="p">)</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
+<span class="w">        </span><span class="sd">&quot;&quot;&quot;</span>
+<span class="sd">        Parameters</span>
+<span class="sd">        ----------</span>
+<span class="sd">        architecture: List[Dict]</span>
+<span class="sd">            List of dictionaries specifying the architecture of the transformer. Each dictionary should have the following keys:</span>
+<span class="sd">            - repeat: int</span>
+<span class="sd">                Number of times to repeat the block</span>
+<span class="sd">            - num_heads: int</span>
+<span class="sd">                Number of heads in the multihead attention</span>
+<span class="sd">            - q_stride: List[int]</span>
+<span class="sd">                Stride for the query in each spatial dimension</span>
+<span class="sd">            - self_attention: bool</span>
+<span class="sd">                Whether to use self attention or mask unit attention</span>
+<span class="sd">        spatial_dims: int</span>
+<span class="sd">            Number of spatial dimensions</span>
+<span class="sd">        num_patches: int, List[int]</span>
+<span class="sd">            Number of patches in each dimension (Z)YX order. If int, the same number of patches is used in each dimension.</span>
+<span class="sd">        num_mask_units: int, List[int]</span>
+<span class="sd">            Number of mask units in each dimension (Z)YX order. If int, the same number of mask units is used in each dimension.</span>
+<span class="sd">        patch_size: int, List[int]</span>
+<span class="sd">            Size of each patch (Z)YX order. If int, the same patch size is used in each dimension.</span>
+<span class="sd">        emb_dim: int</span>
+<span class="sd">            Dimension of embedding</span>
+<span class="sd">        decoder_layer: int</span>
+<span class="sd">            Number of layers in the decoder</span>
+<span class="sd">        decoder_head: int</span>
+<span class="sd">            Number of heads in the decoder</span>
+<span class="sd">        decoder_dim: int</span>
+<span class="sd">            Dimension of the decoder</span>
+<span class="sd">        mask_ratio: float</span>
+<span class="sd">            Fraction of mask units to remove</span>
+<span class="sd">        use_crossmae: bool</span>
+<span class="sd">            Use CrossMAE-style decoder instead of MAE decoder</span>
+<span class="sd">        context_pixels: List[int]</span>
+<span class="sd">            Number of extra pixels around each patch to include in convolutional embedding to encoder dimension.</span>
+<span class="sd">        input_channels: int</span>
+<span class="sd">            Number of input channels</span>
+<span class="sd">        features_only: bool</span>
+<span class="sd">            Only use encoder to extract features</span>
+<span class="sd">        &quot;&quot;&quot;</span>
+        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span>
+            <span class="n">spatial_dims</span><span class="o">=</span><span class="n">spatial_dims</span><span class="p">,</span>
+            <span class="n">num_patches</span><span class="o">=</span><span class="n">num_patches</span><span class="p">,</span>
+            <span class="n">patch_size</span><span class="o">=</span><span class="n">patch_size</span><span class="p">,</span>
+            <span class="n">mask_ratio</span><span class="o">=</span><span class="n">mask_ratio</span><span class="p">,</span>
+            <span class="n">features_only</span><span class="o">=</span><span class="n">features_only</span><span class="p">,</span>
+            <span class="n">context_pixels</span><span class="o">=</span><span class="n">context_pixels</span><span class="p">,</span>
+        <span class="p">)</span>
+        <span class="n">num_mask_units</span> <span class="o">=</span> <span class="n">match_tuple_dimensions</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">spatial_dims</span><span class="p">,</span> <span class="p">[</span><span class="n">num_mask_units</span><span class="p">])[</span><span class="mi">0</span><span class="p">]</span>
+
+        <span class="bp">self</span><span class="o">.</span><span class="n">_encoder</span> <span class="o">=</span> <span class="n">HieraEncoder</span><span class="p">(</span>
+            <span class="n">num_patches</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">num_patches</span><span class="p">,</span>
+            <span class="n">num_mask_units</span><span class="o">=</span><span class="n">num_mask_units</span><span class="p">,</span>
+            <span class="n">architecture</span><span class="o">=</span><span class="n">architecture</span><span class="p">,</span>
+            <span class="n">emb_dim</span><span class="o">=</span><span class="n">emb_dim</span><span class="p">,</span>
+            <span class="n">spatial_dims</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">spatial_dims</span><span class="p">,</span>
+            <span class="n">patch_size</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">patch_size</span><span class="p">,</span>
+            <span class="n">context_pixels</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">context_pixels</span><span class="p">,</span>
+            <span class="n">input_channels</span><span class="o">=</span><span class="n">input_channels</span><span class="p">,</span>
+        <span class="p">)</span>
+        <span class="c1"># &quot;patches&quot; to the decoder are actually mask units, so num_patches is num_mask_units, patch_size is mask unit size</span>
+        <span class="n">mask_unit_size</span> <span class="o">=</span> <span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">num_patches</span><span class="p">)</span> <span class="o">*</span> <span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">patch_size</span><span class="p">))</span> <span class="o">/</span> <span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">(</span>
+            <span class="n">num_mask_units</span>
+        <span class="p">)</span>
+
+        <span class="n">decoder_class</span> <span class="o">=</span> <span class="n">MAE_Decoder</span>
+        <span class="k">if</span> <span class="n">use_crossmae</span><span class="p">:</span>
+            <span class="n">decoder_class</span> <span class="o">=</span> <span class="n">CrossMAE_Decoder</span>
+
+        <span class="bp">self</span><span class="o">.</span><span class="n">_decoder</span> <span class="o">=</span> <span class="n">decoder_class</span><span class="p">(</span>
+            <span class="n">num_patches</span><span class="o">=</span><span class="n">num_mask_units</span><span class="p">,</span>
+            <span class="n">spatial_dims</span><span class="o">=</span><span class="n">spatial_dims</span><span class="p">,</span>
+            <span class="n">patch_size</span><span class="o">=</span><span class="n">mask_unit_size</span><span class="o">.</span><span class="n">astype</span><span class="p">(</span><span class="nb">int</span><span class="p">),</span>
+            <span class="n">enc_dim</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">encoder</span><span class="o">.</span><span class="n">final_dim</span><span class="p">,</span>
+            <span class="n">emb_dim</span><span class="o">=</span><span class="n">decoder_dim</span><span class="p">,</span>
+            <span class="n">num_layer</span><span class="o">=</span><span class="n">decoder_layer</span><span class="p">,</span>
+            <span class="n">num_head</span><span class="o">=</span><span class="n">decoder_head</span><span class="p">,</span>
+            <span class="n">has_cls_token</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
+        <span class="p">)</span>
+
+    <span class="nd">@property</span>
+    <span class="k">def</span> <span class="nf">encoder</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
+        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_encoder</span>
+
+    <span class="nd">@property</span>
+    <span class="k">def</span> <span class="nf">decoder</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
+        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_decoder</span></div>
 </pre></div>
         </article>
       </div>
diff --git a/_modules/cyto_dl/nn/vits/jepa.html b/_modules/cyto_dl/nn/vits/predictor.html
similarity index 83%
rename from _modules/cyto_dl/nn/vits/jepa.html
rename to _modules/cyto_dl/nn/vits/predictor.html
index 2676d76ff..a36a9ce5a 100644
--- a/_modules/cyto_dl/nn/vits/jepa.html
+++ b/_modules/cyto_dl/nn/vits/predictor.html
@@ -5,7 +5,7 @@
     <meta name="color-scheme" content="light dark"><link rel="index" title="Index" href="../../../../genindex.html" /><link rel="search" title="Search" href="../../../../search.html" />
 
     <!-- Generated with Sphinx 6.1.3 and Furo 2023.03.23 -->
-        <title>cyto_dl.nn.vits.jepa - cyto-dl 0.4.0 documentation</title>
+        <title>cyto_dl.nn.vits.predictor - cyto-dl 0.4.0 documentation</title>
       <link rel="stylesheet" type="text/css" href="../../../../_static/pygments.css" />
     <link rel="stylesheet" type="text/css" href="../../../../_static/styles/furo.css?digest=fad236701ea90a88636c2a8c73b44ae642ed2a53" />
     <link rel="stylesheet" type="text/css" href="../../../../_static/styles/furo-extensions.css?digest=30d1aed668e5c3a91c3e3bf6a60b675221979f0e" />
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
@@ -432,102 +439,44 @@
           </label>
         </div>
         <article role="main">
-          <h1>Source code for cyto_dl.nn.vits.jepa</h1><div class="highlight"><pre>
+          <h1>Source code for cyto_dl.nn.vits.predictor</h1><div class="highlight"><pre>
 <span></span><span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">List</span><span class="p">,</span> <span class="n">Optional</span><span class="p">,</span> <span class="n">Union</span>
 
 <span class="kn">import</span> <span class="nn">torch</span>
 <span class="kn">import</span> <span class="nn">torch.nn.functional</span>
 <span class="kn">from</span> <span class="nn">einops</span> <span class="kn">import</span> <span class="n">rearrange</span>
 <span class="kn">from</span> <span class="nn">timm.models.layers</span> <span class="kn">import</span> <span class="n">trunc_normal_</span>
-<span class="kn">from</span> <span class="nn">timm.models.vision_transformer</span> <span class="kn">import</span> <span class="n">Block</span>
 
-<span class="kn">from</span> <span class="nn">cyto_dl.nn.vits.blocks</span> <span class="kn">import</span> <span class="n">CrossAttentionBlock</span><span class="p">,</span> <span class="n">Patchify</span>
-<span class="kn">from</span> <span class="nn">cyto_dl.nn.vits.utils</span> <span class="kn">import</span> <span class="n">get_positional_embedding</span><span class="p">,</span> <span class="n">take_indexes</span>
+<span class="kn">from</span> <span class="nn">cyto_dl.nn.vits.blocks</span> <span class="kn">import</span> <span class="n">CrossAttentionBlock</span>
+<span class="kn">from</span> <span class="nn">cyto_dl.nn.vits.utils</span> <span class="kn">import</span> <span class="p">(</span>
+    <span class="n">get_positional_embedding</span><span class="p">,</span>
+    <span class="n">match_tuple_dimensions</span><span class="p">,</span>
+    <span class="n">take_indexes</span><span class="p">,</span>
+<span class="p">)</span>
+
 
+<div class="viewcode-block" id="JEPAPredictor"><a class="viewcode-back" href="../../../../cyto_dl.nn.vits.predictor.html#cyto_dl.nn.vits.predictor.JEPAPredictor">[docs]</a><span class="k">class</span> <span class="nc">JEPAPredictor</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
+<span class="w">    </span><span class="sd">&quot;&quot;&quot;Class for predicting target features from context embedding.&quot;&quot;&quot;</span>
 
-<div class="viewcode-block" id="JEPAEncoder"><a class="viewcode-back" href="../../../../cyto_dl.nn.vits.jepa.html#cyto_dl.nn.vits.jepa.JEPAEncoder">[docs]</a><span class="k">class</span> <span class="nc">JEPAEncoder</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
     <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span>
         <span class="bp">self</span><span class="p">,</span>
         <span class="n">num_patches</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">]],</span>
         <span class="n">spatial_dims</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">3</span><span class="p">,</span>
-        <span class="n">patch_size</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">]]</span> <span class="o">=</span> <span class="p">(</span><span class="mi">16</span><span class="p">,</span> <span class="mi">16</span><span class="p">,</span> <span class="mi">16</span><span class="p">),</span>
+        <span class="n">input_dim</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">192</span><span class="p">,</span>
         <span class="n">emb_dim</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">192</span><span class="p">,</span>
         <span class="n">num_layer</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">12</span><span class="p">,</span>
         <span class="n">num_head</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">3</span><span class="p">,</span>
-        <span class="n">context_pixels</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">]]</span> <span class="o">=</span> <span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">],</span>
-        <span class="n">input_channels</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">1</span><span class="p">,</span>
         <span class="n">learnable_pos_embedding</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">True</span><span class="p">,</span>
     <span class="p">)</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
 <span class="w">        </span><span class="sd">&quot;&quot;&quot;</span>
 <span class="sd">        Parameters</span>
 <span class="sd">        ----------</span>
-<span class="sd">        num_patches: List[int]</span>
-<span class="sd">            Number of patches in each dimension</span>
+<span class="sd">        num_patches: List[int], int</span>
+<span class="sd">            Number of patches in each dimension. If int, the same number of patches is used for all spatial dimensions</span>
 <span class="sd">        spatial_dims: int</span>
 <span class="sd">            Number of spatial dimensions</span>
-<span class="sd">        patch_size: List[int]</span>
-<span class="sd">            Size of each patch</span>
-<span class="sd">        emb_dim: int</span>
-<span class="sd">            Dimension of embedding</span>
-<span class="sd">        num_layer: int</span>
-<span class="sd">            Number of transformer layers</span>
-<span class="sd">        num_head: int</span>
-<span class="sd">            Number of heads in transformer</span>
-<span class="sd">        context_pixels: List[int]</span>
-<span class="sd">            Number of extra pixels around each patch to include in convolutional embedding to encoder dimension.</span>
-<span class="sd">        input_channels: int</span>
-<span class="sd">            Number of input channels</span>
-<span class="sd">        learnable_pos_embedding: bool</span>
-<span class="sd">            If True, learnable positional embeddings are used. If False, fixed sin/cos positional embeddings. Empirically, fixed positional embeddings work better for brightfield images.</span>
-<span class="sd">        &quot;&quot;&quot;</span>
-        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>
-        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">num_patches</span><span class="p">,</span> <span class="nb">int</span><span class="p">):</span>
-            <span class="n">num_patches</span> <span class="o">=</span> <span class="p">[</span><span class="n">num_patches</span><span class="p">]</span> <span class="o">*</span> <span class="n">spatial_dims</span>
-        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">patch_size</span><span class="p">,</span> <span class="nb">int</span><span class="p">):</span>
-            <span class="n">patch_size</span> <span class="o">=</span> <span class="p">[</span><span class="n">patch_size</span><span class="p">]</span> <span class="o">*</span> <span class="n">spatial_dims</span>
-        <span class="bp">self</span><span class="o">.</span><span class="n">patchify</span> <span class="o">=</span> <span class="n">Patchify</span><span class="p">(</span>
-            <span class="n">patch_size</span><span class="p">,</span>
-            <span class="n">emb_dim</span><span class="p">,</span>
-            <span class="n">num_patches</span><span class="p">,</span>
-            <span class="n">spatial_dims</span><span class="p">,</span>
-            <span class="n">context_pixels</span><span class="p">,</span>
-            <span class="n">input_channels</span><span class="p">,</span>
-            <span class="n">learnable_pos_embedding</span><span class="o">=</span><span class="n">learnable_pos_embedding</span><span class="p">,</span>
-        <span class="p">)</span>
-
-        <span class="bp">self</span><span class="o">.</span><span class="n">transformer</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Sequential</span><span class="p">(</span>
-            <span class="o">*</span><span class="p">[</span><span class="n">Block</span><span class="p">(</span><span class="n">emb_dim</span><span class="p">,</span> <span class="n">num_head</span><span class="p">)</span> <span class="k">for</span> <span class="n">_</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">num_layer</span><span class="p">)]</span>
-        <span class="p">)</span>
-
-        <span class="bp">self</span><span class="o">.</span><span class="n">layer_norm</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">LayerNorm</span><span class="p">(</span><span class="n">emb_dim</span><span class="p">)</span>
-
-<div class="viewcode-block" id="JEPAEncoder.forward"><a class="viewcode-back" href="../../../../cyto_dl.nn.vits.jepa.html#cyto_dl.nn.vits.jepa.JEPAEncoder.forward">[docs]</a>    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">img</span><span class="p">,</span> <span class="n">patchify</span><span class="o">=</span><span class="kc">True</span><span class="p">):</span>
-        <span class="k">if</span> <span class="n">patchify</span><span class="p">:</span>
-            <span class="n">patches</span><span class="p">,</span> <span class="n">_</span><span class="p">,</span> <span class="n">_</span><span class="p">,</span> <span class="n">_</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">patchify</span><span class="p">(</span><span class="n">img</span><span class="p">,</span> <span class="n">mask_ratio</span><span class="o">=</span><span class="mi">0</span><span class="p">)</span>
-        <span class="k">else</span><span class="p">:</span>
-            <span class="n">patches</span> <span class="o">=</span> <span class="n">img</span>
-        <span class="n">patches</span> <span class="o">=</span> <span class="n">rearrange</span><span class="p">(</span><span class="n">patches</span><span class="p">,</span> <span class="s2">&quot;t b c -&gt; b t c&quot;</span><span class="p">)</span>
-        <span class="n">features</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">layer_norm</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">transformer</span><span class="p">(</span><span class="n">patches</span><span class="p">))</span>
-        <span class="k">return</span> <span class="n">features</span></div></div>
-
-
-<div class="viewcode-block" id="JEPAPredictor"><a class="viewcode-back" href="../../../../cyto_dl.nn.vits.jepa.html#cyto_dl.nn.vits.jepa.JEPAPredictor">[docs]</a><span class="k">class</span> <span class="nc">JEPAPredictor</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
-<span class="w">    </span><span class="sd">&quot;&quot;&quot;Class for predicting target features from context embedding.&quot;&quot;&quot;</span>
-
-    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span>
-        <span class="bp">self</span><span class="p">,</span>
-        <span class="n">num_patches</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span>
-        <span class="n">input_dim</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">192</span><span class="p">,</span>
-        <span class="n">emb_dim</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">192</span><span class="p">,</span>
-        <span class="n">num_layer</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">12</span><span class="p">,</span>
-        <span class="n">num_head</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">3</span><span class="p">,</span>
-        <span class="n">learnable_pos_embedding</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">True</span><span class="p">,</span>
-    <span class="p">)</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
-<span class="w">        </span><span class="sd">&quot;&quot;&quot;</span>
-<span class="sd">        Parameters</span>
-<span class="sd">        ----------</span>
-<span class="sd">        num_patches: List[int]</span>
-<span class="sd">            Number of patches in each dimension</span>
+<span class="sd">        input_dim: int</span>
+<span class="sd">            Dimension of input</span>
 <span class="sd">        emb_dim: int</span>
 <span class="sd">            Dimension of embedding</span>
 <span class="sd">        num_layer: int</span>
@@ -549,6 +498,8 @@ <h1>Source code for cyto_dl.nn.vits.jepa</h1><div class="highlight"><pre>
             <span class="p">]</span>
         <span class="p">)</span>
 
+        <span class="n">num_patches</span> <span class="o">=</span> <span class="n">match_tuple_dimensions</span><span class="p">(</span><span class="n">spatial_dims</span><span class="p">,</span> <span class="p">[</span><span class="n">num_patches</span><span class="p">])[</span><span class="mi">0</span><span class="p">]</span>
+
         <span class="bp">self</span><span class="o">.</span><span class="n">mask_token</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Parameter</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="n">emb_dim</span><span class="p">))</span>
         <span class="bp">self</span><span class="o">.</span><span class="n">pos_embedding</span> <span class="o">=</span> <span class="n">get_positional_embedding</span><span class="p">(</span>
             <span class="n">num_patches</span><span class="p">,</span> <span class="n">emb_dim</span><span class="p">,</span> <span class="n">use_cls_token</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">learnable</span><span class="o">=</span><span class="n">learnable_pos_embedding</span>
@@ -560,11 +511,11 @@ <h1>Source code for cyto_dl.nn.vits.jepa</h1><div class="highlight"><pre>
         <span class="bp">self</span><span class="o">.</span><span class="n">norm</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">LayerNorm</span><span class="p">(</span><span class="n">emb_dim</span><span class="p">)</span>
         <span class="bp">self</span><span class="o">.</span><span class="n">init_weight</span><span class="p">()</span>
 
-<div class="viewcode-block" id="JEPAPredictor.init_weight"><a class="viewcode-back" href="../../../../cyto_dl.nn.vits.jepa.html#cyto_dl.nn.vits.jepa.JEPAPredictor.init_weight">[docs]</a>    <span class="k">def</span> <span class="nf">init_weight</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
+<div class="viewcode-block" id="JEPAPredictor.init_weight"><a class="viewcode-back" href="../../../../cyto_dl.nn.vits.predictor.html#cyto_dl.nn.vits.predictor.JEPAPredictor.init_weight">[docs]</a>    <span class="k">def</span> <span class="nf">init_weight</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
         <span class="n">trunc_normal_</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">mask_token</span><span class="p">,</span> <span class="n">std</span><span class="o">=</span><span class="mf">0.02</span><span class="p">)</span>
         <span class="n">trunc_normal_</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">pos_embedding</span><span class="p">,</span> <span class="n">std</span><span class="o">=</span><span class="mf">0.02</span><span class="p">)</span></div>
 
-<div class="viewcode-block" id="JEPAPredictor.predict_target_features"><a class="viewcode-back" href="../../../../cyto_dl.nn.vits.jepa.html#cyto_dl.nn.vits.jepa.JEPAPredictor.predict_target_features">[docs]</a>    <span class="k">def</span> <span class="nf">predict_target_features</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">context_emb</span><span class="p">,</span> <span class="n">target_masks</span><span class="p">):</span>
+<div class="viewcode-block" id="JEPAPredictor.predict_target_features"><a class="viewcode-back" href="../../../../cyto_dl.nn.vits.predictor.html#cyto_dl.nn.vits.predictor.JEPAPredictor.predict_target_features">[docs]</a>    <span class="k">def</span> <span class="nf">predict_target_features</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">context_emb</span><span class="p">,</span> <span class="n">target_masks</span><span class="p">):</span>
         <span class="n">t</span><span class="p">,</span> <span class="n">b</span> <span class="o">=</span> <span class="n">target_masks</span><span class="o">.</span><span class="n">shape</span>
         <span class="c1"># add masked positional embedding to mask tokens</span>
         <span class="n">mask</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">mask_token</span><span class="o">.</span><span class="n">expand</span><span class="p">(</span><span class="n">t</span><span class="p">,</span> <span class="n">b</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">)</span>
@@ -581,21 +532,22 @@ <h1>Source code for cyto_dl.nn.vits.jepa</h1><div class="highlight"><pre>
         <span class="n">mask</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">projector_embed</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">norm</span><span class="p">(</span><span class="n">mask</span><span class="p">))</span>
         <span class="k">return</span> <span class="n">mask</span></div>
 
-<div class="viewcode-block" id="JEPAPredictor.forward"><a class="viewcode-back" href="../../../../cyto_dl.nn.vits.jepa.html#cyto_dl.nn.vits.jepa.JEPAPredictor.forward">[docs]</a>    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">context_emb</span><span class="p">,</span> <span class="n">target_masks</span><span class="p">):</span>
+<div class="viewcode-block" id="JEPAPredictor.forward"><a class="viewcode-back" href="../../../../cyto_dl.nn.vits.predictor.html#cyto_dl.nn.vits.predictor.JEPAPredictor.forward">[docs]</a>    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">context_emb</span><span class="p">,</span> <span class="n">target_masks</span><span class="p">):</span>
         <span class="c1"># map context embedding to predictor dimension</span>
         <span class="n">context_emb</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">predictor_embed</span><span class="p">(</span><span class="n">context_emb</span><span class="p">)</span>
         <span class="n">target_features</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">predict_target_features</span><span class="p">(</span><span class="n">context_emb</span><span class="p">,</span> <span class="n">target_masks</span><span class="p">)</span>
         <span class="k">return</span> <span class="n">target_features</span></div></div>
 
 
-<div class="viewcode-block" id="IWMPredictor"><a class="viewcode-back" href="../../../../cyto_dl.nn.vits.jepa.html#cyto_dl.nn.vits.jepa.IWMPredictor">[docs]</a><span class="k">class</span> <span class="nc">IWMPredictor</span><span class="p">(</span><span class="n">JEPAPredictor</span><span class="p">):</span>
+<div class="viewcode-block" id="IWMPredictor"><a class="viewcode-back" href="../../../../cyto_dl.nn.vits.predictor.html#cyto_dl.nn.vits.predictor.IWMPredictor">[docs]</a><span class="k">class</span> <span class="nc">IWMPredictor</span><span class="p">(</span><span class="n">JEPAPredictor</span><span class="p">):</span>
 <span class="w">    </span><span class="sd">&quot;&quot;&quot;Specialized JEPA predictor that can conditionally predict between different domains (e.g.</span>
 <span class="sd">    predict from brightfield to multiple fluorescent tags)&quot;&quot;&quot;</span>
 
     <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span>
         <span class="bp">self</span><span class="p">,</span>
         <span class="n">domains</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="nb">str</span><span class="p">],</span>
-        <span class="n">num_patches</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span>
+        <span class="n">num_patches</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">]],</span>
+        <span class="n">spatial_dims</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">3</span><span class="p">,</span>
         <span class="n">input_dim</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">192</span><span class="p">,</span>
         <span class="n">emb_dim</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">192</span><span class="p">,</span>
         <span class="n">num_layer</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">12</span><span class="p">,</span>
@@ -607,7 +559,11 @@ <h1>Source code for cyto_dl.nn.vits.jepa</h1><div class="highlight"><pre>
 <span class="sd">        domains: List[str]</span>
 <span class="sd">            List of names of target domains</span>
 <span class="sd">        num_patches: List[int]</span>
-<span class="sd">            Number of patches in each dimension</span>
+<span class="sd">            Number of patches in each dimension. If int, the same number of patches is used for all spatial dimensions</span>
+<span class="sd">        spatial_dims: int</span>
+<span class="sd">            Number of spatial dimensions</span>
+<span class="sd">        spatial_dims: int</span>
+<span class="sd">            Number of spatial dimensions</span>
 <span class="sd">        emb_dim: int</span>
 <span class="sd">            Dimension of embedding</span>
 <span class="sd">        num_layer: int</span>
@@ -615,14 +571,21 @@ <h1>Source code for cyto_dl.nn.vits.jepa</h1><div class="highlight"><pre>
 <span class="sd">        num_head: int</span>
 <span class="sd">            Number of heads in transformer</span>
 <span class="sd">        &quot;&quot;&quot;</span>
-        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">num_patches</span><span class="p">,</span> <span class="n">input_dim</span><span class="p">,</span> <span class="n">emb_dim</span><span class="p">,</span> <span class="n">num_layer</span><span class="p">,</span> <span class="n">num_head</span><span class="p">)</span>
+        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span>
+            <span class="n">num_patches</span><span class="o">=</span><span class="n">num_patches</span><span class="p">,</span>
+            <span class="n">spatial_dims</span><span class="o">=</span><span class="n">spatial_dims</span><span class="p">,</span>
+            <span class="n">input_dim</span><span class="o">=</span><span class="n">input_dim</span><span class="p">,</span>
+            <span class="n">emb_dim</span><span class="o">=</span><span class="n">emb_dim</span><span class="p">,</span>
+            <span class="n">num_layer</span><span class="o">=</span><span class="n">num_layer</span><span class="p">,</span>
+            <span class="n">num_head</span><span class="o">=</span><span class="n">num_head</span><span class="p">,</span>
+        <span class="p">)</span>
 
         <span class="bp">self</span><span class="o">.</span><span class="n">domain_embeddings</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">ParameterDict</span><span class="p">(</span>
             <span class="p">{</span><span class="n">d</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Parameter</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="n">emb_dim</span><span class="p">))</span> <span class="k">for</span> <span class="n">d</span> <span class="ow">in</span> <span class="n">domains</span><span class="p">}</span>
         <span class="p">)</span>
         <span class="bp">self</span><span class="o">.</span><span class="n">context_mixer</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Linear</span><span class="p">(</span><span class="mi">2</span> <span class="o">*</span> <span class="n">emb_dim</span><span class="p">,</span> <span class="n">emb_dim</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span>
 
-<div class="viewcode-block" id="IWMPredictor.forward"><a class="viewcode-back" href="../../../../cyto_dl.nn.vits.jepa.html#cyto_dl.nn.vits.jepa.IWMPredictor.forward">[docs]</a>    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">context_emb</span><span class="p">,</span> <span class="n">target_masks</span><span class="p">,</span> <span class="n">target_domain</span><span class="p">):</span>
+<div class="viewcode-block" id="IWMPredictor.forward"><a class="viewcode-back" href="../../../../cyto_dl.nn.vits.predictor.html#cyto_dl.nn.vits.predictor.IWMPredictor.forward">[docs]</a>    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">context_emb</span><span class="p">,</span> <span class="n">target_masks</span><span class="p">,</span> <span class="n">target_domain</span><span class="p">):</span>
         <span class="n">_</span><span class="p">,</span> <span class="n">b</span> <span class="o">=</span> <span class="n">target_masks</span><span class="o">.</span><span class="n">shape</span>
         <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">target_domain</span><span class="p">)</span> <span class="o">==</span> <span class="mi">1</span><span class="p">:</span>
             <span class="n">target_domain</span> <span class="o">=</span> <span class="n">target_domain</span> <span class="o">*</span> <span class="n">b</span>
diff --git a/_modules/cyto_dl/nn/vits/seg.html b/_modules/cyto_dl/nn/vits/seg.html
index d62af0482..c1f12d256 100644
--- a/_modules/cyto_dl/nn/vits/seg.html
+++ b/_modules/cyto_dl/nn/vits/seg.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/nn/vits/utils.html b/_modules/cyto_dl/nn/vits/utils.html
index 2286412fa..bb2ed3f75 100644
--- a/_modules/cyto_dl/nn/vits/utils.html
+++ b/_modules/cyto_dl/nn/vits/utils.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
@@ -438,6 +445,7 @@ <h1>Source code for cyto_dl.nn.vits.utils</h1><div class="highlight"><pre>
 <span class="kn">import</span> <span class="nn">numpy</span> <span class="k">as</span> <span class="nn">np</span>
 <span class="kn">import</span> <span class="nn">torch</span>
 <span class="kn">from</span> <span class="nn">einops</span> <span class="kn">import</span> <span class="n">rearrange</span><span class="p">,</span> <span class="n">repeat</span>
+<span class="kn">from</span> <span class="nn">monai.utils.misc</span> <span class="kn">import</span> <span class="n">ensure_tuple_rep</span>
 <span class="kn">from</span> <span class="nn">positional_encodings.torch_encodings</span> <span class="kn">import</span> <span class="p">(</span>
     <span class="n">PositionalEncoding2D</span><span class="p">,</span>
     <span class="n">PositionalEncoding3D</span><span class="p">,</span>
@@ -449,6 +457,12 @@ <h1>Source code for cyto_dl.nn.vits.utils</h1><div class="highlight"><pre>
     <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">gather</span><span class="p">(</span><span class="n">sequences</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="n">repeat</span><span class="p">(</span><span class="n">indexes</span><span class="p">,</span> <span class="s2">&quot;t b -&gt; t b c&quot;</span><span class="p">,</span> <span class="n">c</span><span class="o">=</span><span class="n">sequences</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">]))</span></div>
 
 
+<div class="viewcode-block" id="random_indexes"><a class="viewcode-back" href="../../../../cyto_dl.nn.vits.utils.html#cyto_dl.nn.vits.utils.random_indexes">[docs]</a><span class="k">def</span> <span class="nf">random_indexes</span><span class="p">(</span><span class="n">size</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">device</span><span class="p">):</span>
+    <span class="n">forward_indexes</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">randperm</span><span class="p">(</span><span class="n">size</span><span class="p">,</span> <span class="n">device</span><span class="o">=</span><span class="n">device</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">torch</span><span class="o">.</span><span class="n">long</span><span class="p">)</span>
+    <span class="n">backward_indexes</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">argsort</span><span class="p">(</span><span class="n">forward_indexes</span><span class="p">)</span>
+    <span class="k">return</span> <span class="n">forward_indexes</span><span class="p">,</span> <span class="n">backward_indexes</span></div>
+
+
 <div class="viewcode-block" id="get_positional_embedding"><a class="viewcode-back" href="../../../../cyto_dl.nn.vits.utils.html#cyto_dl.nn.vits.utils.get_positional_embedding">[docs]</a><span class="k">def</span> <span class="nf">get_positional_embedding</span><span class="p">(</span>
     <span class="n">num_patches</span><span class="p">:</span> <span class="n">Sequence</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span> <span class="n">emb_dim</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">use_cls_token</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">True</span><span class="p">,</span> <span class="n">learnable</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">True</span>
 <span class="p">):</span>
@@ -475,6 +489,15 @@ <h1>Source code for cyto_dl.nn.vits.utils</h1><div class="highlight"><pre>
             <span class="n">cls_token</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="n">emb_dim</span><span class="p">)</span>
             <span class="n">pe</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">cat</span><span class="p">([</span><span class="n">cls_token</span><span class="p">,</span> <span class="n">pe</span><span class="p">],</span> <span class="n">dim</span><span class="o">=</span><span class="mi">0</span><span class="p">)</span>
         <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Parameter</span><span class="p">(</span><span class="n">pe</span><span class="p">,</span> <span class="n">requires_grad</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span></div>
+
+
+<div class="viewcode-block" id="match_tuple_dimensions"><a class="viewcode-back" href="../../../../cyto_dl.nn.vits.utils.html#cyto_dl.nn.vits.utils.match_tuple_dimensions">[docs]</a><span class="k">def</span> <span class="nf">match_tuple_dimensions</span><span class="p">(</span><span class="n">spatial_dims</span><span class="p">,</span> <span class="n">tuples</span><span class="p">):</span>
+<span class="w">    </span><span class="sd">&quot;&quot;&quot;Ensure that each element in a list of tuples has the same length as spatial_dims.</span>
+
+<span class="sd">    If a single element, the element is repeated to match the spatial_dims.</span>
+<span class="sd">    &quot;&quot;&quot;</span>
+    <span class="k">assert</span> <span class="n">spatial_dims</span> <span class="ow">in</span> <span class="p">(</span><span class="mi">2</span><span class="p">,</span> <span class="mi">3</span><span class="p">),</span> <span class="s2">&quot;spatial_dims must be 2 or 3&quot;</span>
+    <span class="k">return</span> <span class="p">[</span><span class="n">ensure_tuple_rep</span><span class="p">(</span><span class="n">t</span><span class="p">,</span> <span class="n">spatial_dims</span><span class="p">)</span> <span class="k">for</span> <span class="n">t</span> <span class="ow">in</span> <span class="n">tuples</span><span class="p">]</span></div>
 </pre></div>
         </article>
       </div>
diff --git a/_modules/cyto_dl/train.html b/_modules/cyto_dl/train.html
index c1cb447a5..e54024bc5 100644
--- a/_modules/cyto_dl/train.html
+++ b/_modules/cyto_dl/train.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/utils/array.html b/_modules/cyto_dl/utils/array.html
index cdb062960..9bc9fff3f 100644
--- a/_modules/cyto_dl/utils/array.html
+++ b/_modules/cyto_dl/utils/array.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/utils/checkpoint.html b/_modules/cyto_dl/utils/checkpoint.html
index 7ff744286..bc622c643 100644
--- a/_modules/cyto_dl/utils/checkpoint.html
+++ b/_modules/cyto_dl/utils/checkpoint.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/utils/config.html b/_modules/cyto_dl/utils/config.html
index 08fe0d761..5d2f3bcab 100644
--- a/_modules/cyto_dl/utils/config.html
+++ b/_modules/cyto_dl/utils/config.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/utils/download_test_data.html b/_modules/cyto_dl/utils/download_test_data.html
index a3c870627..bae33c029 100644
--- a/_modules/cyto_dl/utils/download_test_data.html
+++ b/_modules/cyto_dl/utils/download_test_data.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/utils/dummy_dataset.html b/_modules/cyto_dl/utils/dummy_dataset.html
index 561cbf781..b159c0d40 100644
--- a/_modules/cyto_dl/utils/dummy_dataset.html
+++ b/_modules/cyto_dl/utils/dummy_dataset.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/utils/pylogger.html b/_modules/cyto_dl/utils/pylogger.html
index d0dc9d59d..9a5aea838 100644
--- a/_modules/cyto_dl/utils/pylogger.html
+++ b/_modules/cyto_dl/utils/pylogger.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/utils/rich_utils.html b/_modules/cyto_dl/utils/rich_utils.html
index 0fa8b798b..70bfc2721 100644
--- a/_modules/cyto_dl/utils/rich_utils.html
+++ b/_modules/cyto_dl/utils/rich_utils.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/utils/rotation.html b/_modules/cyto_dl/utils/rotation.html
index b63f723eb..1cfd8d565 100644
--- a/_modules/cyto_dl/utils/rotation.html
+++ b/_modules/cyto_dl/utils/rotation.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/cyto_dl/utils/template_utils.html b/_modules/cyto_dl/utils/template_utils.html
index e8dc54c68..9e9f9268b 100644
--- a/_modules/cyto_dl/utils/template_utils.html
+++ b/_modules/cyto_dl/utils/template_utils.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../../../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../../../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../../../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../../../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../../../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/_modules/index.html b/_modules/index.html
index b19dfe029..2b62dacd5 100644
--- a/_modules/index.html
+++ b/_modules/index.html
@@ -349,14 +349,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="../cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="../cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="../cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="../cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="../cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="../cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="../cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="../cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="../cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="../cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="../cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -368,15 +375,15 @@
 <li class="toctree-l4"><a class="reference internal" href="../cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="../cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="../cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="../cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="../cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
@@ -523,10 +530,14 @@ <h1>All modules for which code is available</h1>
 <li><a href="cyto_dl/nn/track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor</a></li>
 <li><a href="cyto_dl/nn/vits/blocks/cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention</a></li>
 <li><a href="cyto_dl/nn/vits/blocks/intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher</a></li>
-<li><a href="cyto_dl/nn/vits/blocks/patchify.html">cyto_dl.nn.vits.blocks.patchify</a></li>
-<li><a href="cyto_dl/nn/vits/cross_mae.html">cyto_dl.nn.vits.cross_mae</a></li>
-<li><a href="cyto_dl/nn/vits/jepa.html">cyto_dl.nn.vits.jepa</a></li>
+<li><a href="cyto_dl/nn/vits/blocks/masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention</a></li>
+<li><a href="cyto_dl/nn/vits/blocks/patchify/patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify</a></li>
+<li><a href="cyto_dl/nn/vits/blocks/patchify/patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base</a></li>
+<li><a href="cyto_dl/nn/vits/blocks/patchify/patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera</a></li>
+<li><a href="cyto_dl/nn/vits/decoder.html">cyto_dl.nn.vits.decoder</a></li>
+<li><a href="cyto_dl/nn/vits/encoder.html">cyto_dl.nn.vits.encoder</a></li>
 <li><a href="cyto_dl/nn/vits/mae.html">cyto_dl.nn.vits.mae</a></li>
+<li><a href="cyto_dl/nn/vits/predictor.html">cyto_dl.nn.vits.predictor</a></li>
 <li><a href="cyto_dl/nn/vits/seg.html">cyto_dl.nn.vits.seg</a></li>
 <li><a href="cyto_dl/nn/vits/utils.html">cyto_dl.nn.vits.utils</a></li>
 <li><a href="cyto_dl/train.html">cyto_dl.train</a></li>
diff --git a/_sources/cyto_dl.nn.vits.blocks.masked_unit_attention.rst.txt b/_sources/cyto_dl.nn.vits.blocks.masked_unit_attention.rst.txt
new file mode 100644
index 000000000..8339efb2d
--- /dev/null
+++ b/_sources/cyto_dl.nn.vits.blocks.masked_unit_attention.rst.txt
@@ -0,0 +1,7 @@
+cyto\_dl.nn.vits.blocks.masked\_unit\_attention module
+======================================================
+
+.. automodule:: cyto_dl.nn.vits.blocks.masked_unit_attention
+   :members:
+   :undoc-members:
+   :show-inheritance:
diff --git a/_sources/cyto_dl.nn.vits.blocks.patchify.patchify.rst.txt b/_sources/cyto_dl.nn.vits.blocks.patchify.patchify.rst.txt
new file mode 100644
index 000000000..c33f201c4
--- /dev/null
+++ b/_sources/cyto_dl.nn.vits.blocks.patchify.patchify.rst.txt
@@ -0,0 +1,7 @@
+cyto\_dl.nn.vits.blocks.patchify.patchify module
+================================================
+
+.. automodule:: cyto_dl.nn.vits.blocks.patchify.patchify
+   :members:
+   :undoc-members:
+   :show-inheritance:
diff --git a/_sources/cyto_dl.nn.vits.blocks.patchify.patchify_base.rst.txt b/_sources/cyto_dl.nn.vits.blocks.patchify.patchify_base.rst.txt
new file mode 100644
index 000000000..d37f05d2f
--- /dev/null
+++ b/_sources/cyto_dl.nn.vits.blocks.patchify.patchify_base.rst.txt
@@ -0,0 +1,7 @@
+cyto\_dl.nn.vits.blocks.patchify.patchify\_base module
+======================================================
+
+.. automodule:: cyto_dl.nn.vits.blocks.patchify.patchify_base
+   :members:
+   :undoc-members:
+   :show-inheritance:
diff --git a/_sources/cyto_dl.nn.vits.blocks.patchify.patchify_hiera.rst.txt b/_sources/cyto_dl.nn.vits.blocks.patchify.patchify_hiera.rst.txt
new file mode 100644
index 000000000..03a06669b
--- /dev/null
+++ b/_sources/cyto_dl.nn.vits.blocks.patchify.patchify_hiera.rst.txt
@@ -0,0 +1,7 @@
+cyto\_dl.nn.vits.blocks.patchify.patchify\_hiera module
+=======================================================
+
+.. automodule:: cyto_dl.nn.vits.blocks.patchify.patchify_hiera
+   :members:
+   :undoc-members:
+   :show-inheritance:
diff --git a/_sources/cyto_dl.nn.vits.blocks.patchify.rst.txt b/_sources/cyto_dl.nn.vits.blocks.patchify.rst.txt
index 9f6ecee03..9886f1f8b 100644
--- a/_sources/cyto_dl.nn.vits.blocks.patchify.rst.txt
+++ b/_sources/cyto_dl.nn.vits.blocks.patchify.rst.txt
@@ -1,7 +1,17 @@
-cyto\_dl.nn.vits.blocks.patchify module
-=======================================
+cyto\_dl.nn.vits.blocks.patchify package
+========================================
 
 .. automodule:: cyto_dl.nn.vits.blocks.patchify
    :members:
    :undoc-members:
    :show-inheritance:
+
+Submodules
+----------
+
+.. toctree::
+   :maxdepth: 4
+
+   cyto_dl.nn.vits.blocks.patchify.patchify
+   cyto_dl.nn.vits.blocks.patchify.patchify_base
+   cyto_dl.nn.vits.blocks.patchify.patchify_hiera
diff --git a/_sources/cyto_dl.nn.vits.blocks.rst.txt b/_sources/cyto_dl.nn.vits.blocks.rst.txt
index e05283e3e..5a7d05b78 100644
--- a/_sources/cyto_dl.nn.vits.blocks.rst.txt
+++ b/_sources/cyto_dl.nn.vits.blocks.rst.txt
@@ -6,6 +6,14 @@ cyto\_dl.nn.vits.blocks package
    :undoc-members:
    :show-inheritance:
 
+Subpackages
+-----------
+
+.. toctree::
+   :maxdepth: 4
+
+   cyto_dl.nn.vits.blocks.patchify
+
 Submodules
 ----------
 
@@ -14,4 +22,4 @@ Submodules
 
    cyto_dl.nn.vits.blocks.cross_attention
    cyto_dl.nn.vits.blocks.intermediate_weigher
-   cyto_dl.nn.vits.blocks.patchify
+   cyto_dl.nn.vits.blocks.masked_unit_attention
diff --git a/_sources/cyto_dl.nn.vits.cross_mae.rst.txt b/_sources/cyto_dl.nn.vits.cross_mae.rst.txt
deleted file mode 100644
index f8af018b5..000000000
--- a/_sources/cyto_dl.nn.vits.cross_mae.rst.txt
+++ /dev/null
@@ -1,7 +0,0 @@
-cyto\_dl.nn.vits.cross\_mae module
-==================================
-
-.. automodule:: cyto_dl.nn.vits.cross_mae
-   :members:
-   :undoc-members:
-   :show-inheritance:
diff --git a/_sources/cyto_dl.nn.vits.decoder.rst.txt b/_sources/cyto_dl.nn.vits.decoder.rst.txt
new file mode 100644
index 000000000..9e248c9b5
--- /dev/null
+++ b/_sources/cyto_dl.nn.vits.decoder.rst.txt
@@ -0,0 +1,7 @@
+cyto\_dl.nn.vits.decoder module
+===============================
+
+.. automodule:: cyto_dl.nn.vits.decoder
+   :members:
+   :undoc-members:
+   :show-inheritance:
diff --git a/_sources/cyto_dl.nn.vits.encoder.rst.txt b/_sources/cyto_dl.nn.vits.encoder.rst.txt
new file mode 100644
index 000000000..dd1aa23cf
--- /dev/null
+++ b/_sources/cyto_dl.nn.vits.encoder.rst.txt
@@ -0,0 +1,7 @@
+cyto\_dl.nn.vits.encoder module
+===============================
+
+.. automodule:: cyto_dl.nn.vits.encoder
+   :members:
+   :undoc-members:
+   :show-inheritance:
diff --git a/_sources/cyto_dl.nn.vits.jepa.rst.txt b/_sources/cyto_dl.nn.vits.jepa.rst.txt
deleted file mode 100644
index a9336e7db..000000000
--- a/_sources/cyto_dl.nn.vits.jepa.rst.txt
+++ /dev/null
@@ -1,7 +0,0 @@
-cyto\_dl.nn.vits.jepa module
-============================
-
-.. automodule:: cyto_dl.nn.vits.jepa
-   :members:
-   :undoc-members:
-   :show-inheritance:
diff --git a/_sources/cyto_dl.nn.vits.predictor.rst.txt b/_sources/cyto_dl.nn.vits.predictor.rst.txt
new file mode 100644
index 000000000..27fd24916
--- /dev/null
+++ b/_sources/cyto_dl.nn.vits.predictor.rst.txt
@@ -0,0 +1,7 @@
+cyto\_dl.nn.vits.predictor module
+=================================
+
+.. automodule:: cyto_dl.nn.vits.predictor
+   :members:
+   :undoc-members:
+   :show-inheritance:
diff --git a/_sources/cyto_dl.nn.vits.rst.txt b/_sources/cyto_dl.nn.vits.rst.txt
index 9082f8661..44a1bc2ba 100644
--- a/_sources/cyto_dl.nn.vits.rst.txt
+++ b/_sources/cyto_dl.nn.vits.rst.txt
@@ -20,8 +20,9 @@ Submodules
 .. toctree::
    :maxdepth: 4
 
-   cyto_dl.nn.vits.cross_mae
-   cyto_dl.nn.vits.jepa
+   cyto_dl.nn.vits.decoder
+   cyto_dl.nn.vits.encoder
    cyto_dl.nn.vits.mae
+   cyto_dl.nn.vits.predictor
    cyto_dl.nn.vits.seg
    cyto_dl.nn.vits.utils
diff --git a/api.html b/api.html
index 6d46a5d8f..472a9a472 100644
--- a/api.html
+++ b/api.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.api.cyto_dl_model.cyto_dl_base_model.html b/cyto_dl.api.cyto_dl_model.cyto_dl_base_model.html
index bb516cf67..b8554363a 100644
--- a/cyto_dl.api.cyto_dl_model.cyto_dl_base_model.html
+++ b/cyto_dl.api.cyto_dl_model.cyto_dl_base_model.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.api.cyto_dl_model.html b/cyto_dl.api.cyto_dl_model.html
index 1edbd7e07..74c0f7d7c 100644
--- a/cyto_dl.api.cyto_dl_model.html
+++ b/cyto_dl.api.cyto_dl_model.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.api.cyto_dl_model.segmentation_plugin_model.html b/cyto_dl.api.cyto_dl_model.segmentation_plugin_model.html
index ff697e2ec..31750fb26 100644
--- a/cyto_dl.api.cyto_dl_model.segmentation_plugin_model.html
+++ b/cyto_dl.api.cyto_dl_model.segmentation_plugin_model.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.api.data.html b/cyto_dl.api.data.html
index fd2f4db5b..46e27cae0 100644
--- a/cyto_dl.api.data.html
+++ b/cyto_dl.api.data.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.api.html b/cyto_dl.api.html
index 1241e1202..5f2f143d4 100644
--- a/cyto_dl.api.html
+++ b/cyto_dl.api.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.api.model.html b/cyto_dl.api.model.html
index 83aaac6a3..6eb5c91b8 100644
--- a/cyto_dl.api.model.html
+++ b/cyto_dl.api.model.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.callbacks.callback_utils.html b/cyto_dl.callbacks.callback_utils.html
index 04853d4ee..967d76b66 100644
--- a/cyto_dl.callbacks.callback_utils.html
+++ b/cyto_dl.callbacks.callback_utils.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.callbacks.csv_saver.html b/cyto_dl.callbacks.csv_saver.html
index 0679321c5..da1922a28 100644
--- a/cyto_dl.callbacks.csv_saver.html
+++ b/cyto_dl.callbacks.csv_saver.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.callbacks.html b/cyto_dl.callbacks.html
index 40210102a..5503ddccf 100644
--- a/cyto_dl.callbacks.html
+++ b/cyto_dl.callbacks.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.callbacks.image_saver.html b/cyto_dl.callbacks.image_saver.html
index 1c42d82fa..94b271ca7 100644
--- a/cyto_dl.callbacks.image_saver.html
+++ b/cyto_dl.callbacks.image_saver.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.callbacks.latent_walk.html b/cyto_dl.callbacks.latent_walk.html
index b2a649173..2ea5fa886 100644
--- a/cyto_dl.callbacks.latent_walk.html
+++ b/cyto_dl.callbacks.latent_walk.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.callbacks.layer_freeze.html b/cyto_dl.callbacks.layer_freeze.html
index 1444994a0..99b96f643 100644
--- a/cyto_dl.callbacks.layer_freeze.html
+++ b/cyto_dl.callbacks.layer_freeze.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.callbacks.model_utils.html b/cyto_dl.callbacks.model_utils.html
index 47e9f9816..d5a664401 100644
--- a/cyto_dl.callbacks.model_utils.html
+++ b/cyto_dl.callbacks.model_utils.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.callbacks.outlier_detection.html b/cyto_dl.callbacks.outlier_detection.html
index 9666f9282..de5054a9b 100644
--- a/cyto_dl.callbacks.outlier_detection.html
+++ b/cyto_dl.callbacks.outlier_detection.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.compile.html b/cyto_dl.compile.html
index d83cb91bd..9120f4060 100644
--- a/cyto_dl.compile.html
+++ b/cyto_dl.compile.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.dataframe.html b/cyto_dl.dataframe.html
index a0f01b9cc..11f02bd71 100644
--- a/cyto_dl.dataframe.html
+++ b/cyto_dl.dataframe.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.dataframe.readers.html b/cyto_dl.dataframe.readers.html
index 3464f17e1..7e6e19c8b 100644
--- a/cyto_dl.dataframe.readers.html
+++ b/cyto_dl.dataframe.readers.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.dataframe.transforms.filter.html b/cyto_dl.dataframe.transforms.filter.html
index 8ddfd42fc..e91a1c167 100644
--- a/cyto_dl.dataframe.transforms.filter.html
+++ b/cyto_dl.dataframe.transforms.filter.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.dataframe.transforms.group_cols.html b/cyto_dl.dataframe.transforms.group_cols.html
index a2c1efcdc..06cde473f 100644
--- a/cyto_dl.dataframe.transforms.group_cols.html
+++ b/cyto_dl.dataframe.transforms.group_cols.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.dataframe.transforms.html b/cyto_dl.dataframe.transforms.html
index 37fbcb465..23af8fd14 100644
--- a/cyto_dl.dataframe.transforms.html
+++ b/cyto_dl.dataframe.transforms.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.dataframe.transforms.misc.html b/cyto_dl.dataframe.transforms.misc.html
index ce33e9251..9a5cb62e6 100644
--- a/cyto_dl.dataframe.transforms.misc.html
+++ b/cyto_dl.dataframe.transforms.misc.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.dataframe.transforms.split.html b/cyto_dl.dataframe.transforms.split.html
index 9ec0d2e63..73b8728ce 100644
--- a/cyto_dl.dataframe.transforms.split.html
+++ b/cyto_dl.dataframe.transforms.split.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.datamodules.array.html b/cyto_dl.datamodules.array.html
index 35ce89103..be0406f1a 100644
--- a/cyto_dl.datamodules.array.html
+++ b/cyto_dl.datamodules.array.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.datamodules.data_dict.html b/cyto_dl.datamodules.data_dict.html
index 3e9439957..f185aa6e3 100644
--- a/cyto_dl.datamodules.data_dict.html
+++ b/cyto_dl.datamodules.data_dict.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.datamodules.dataframe.dataframe_datamodule.html b/cyto_dl.datamodules.dataframe.dataframe_datamodule.html
index f7d95e3c3..ebd399dc0 100644
--- a/cyto_dl.datamodules.dataframe.dataframe_datamodule.html
+++ b/cyto_dl.datamodules.dataframe.dataframe_datamodule.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.datamodules.dataframe.grouped_dataframe_datamodule.html b/cyto_dl.datamodules.dataframe.grouped_dataframe_datamodule.html
index 4ce9d773c..58f50b68c 100644
--- a/cyto_dl.datamodules.dataframe.grouped_dataframe_datamodule.html
+++ b/cyto_dl.datamodules.dataframe.grouped_dataframe_datamodule.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.datamodules.dataframe.html b/cyto_dl.datamodules.dataframe.html
index 5760cefc4..8a646942f 100644
--- a/cyto_dl.datamodules.dataframe.html
+++ b/cyto_dl.datamodules.dataframe.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.datamodules.dataframe.utils.html b/cyto_dl.datamodules.dataframe.utils.html
index b169f2eb4..94f600f76 100644
--- a/cyto_dl.datamodules.dataframe.utils.html
+++ b/cyto_dl.datamodules.dataframe.utils.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.datamodules.folder.html b/cyto_dl.datamodules.folder.html
index 0d6538df9..88b42ab82 100644
--- a/cyto_dl.datamodules.folder.html
+++ b/cyto_dl.datamodules.folder.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.datamodules.html b/cyto_dl.datamodules.html
index 720335db0..39a6ff467 100644
--- a/cyto_dl.datamodules.html
+++ b/cyto_dl.datamodules.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.datamodules.multidim_image.html b/cyto_dl.datamodules.multidim_image.html
index d24e83e79..1db57df8a 100644
--- a/cyto_dl.datamodules.multidim_image.html
+++ b/cyto_dl.datamodules.multidim_image.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.datamodules.smartcache.html b/cyto_dl.datamodules.smartcache.html
index 971b0a0f1..281c589cb 100644
--- a/cyto_dl.datamodules.smartcache.html
+++ b/cyto_dl.datamodules.smartcache.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.datamodules.torchvision.html b/cyto_dl.datamodules.torchvision.html
index 1404c3c25..be861fd77 100644
--- a/cyto_dl.datamodules.torchvision.html
+++ b/cyto_dl.datamodules.torchvision.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.eval.html b/cyto_dl.eval.html
index 3dd5d8192..af3439867 100644
--- a/cyto_dl.eval.html
+++ b/cyto_dl.eval.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.html b/cyto_dl.html
index 49e9ad514..a0e493dd0 100644
--- a/cyto_dl.html
+++ b/cyto_dl.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.image.html b/cyto_dl.image.html
index 7fba8e8bb..3727bf75e 100644
--- a/cyto_dl.image.html
+++ b/cyto_dl.image.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.image.io.aicsimage_loader.html b/cyto_dl.image.io.aicsimage_loader.html
index 49bbde39b..b69c8def8 100644
--- a/cyto_dl.image.io.aicsimage_loader.html
+++ b/cyto_dl.image.io.aicsimage_loader.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.image.io.html b/cyto_dl.image.io.html
index c5c9a8eb3..ea8c7ff63 100644
--- a/cyto_dl.image.io.html
+++ b/cyto_dl.image.io.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.image.io.monai_bio_reader.html b/cyto_dl.image.io.monai_bio_reader.html
index fb7a5dc9c..3b601d728 100644
--- a/cyto_dl.image.io.monai_bio_reader.html
+++ b/cyto_dl.image.io.monai_bio_reader.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.image.io.numpy_reader.html b/cyto_dl.image.io.numpy_reader.html
index 4c38cbbe2..19ceaf57e 100644
--- a/cyto_dl.image.io.numpy_reader.html
+++ b/cyto_dl.image.io.numpy_reader.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.image.io.ome_zarr_reader.html b/cyto_dl.image.io.ome_zarr_reader.html
index dbc92f492..da77c95a5 100644
--- a/cyto_dl.image.io.ome_zarr_reader.html
+++ b/cyto_dl.image.io.ome_zarr_reader.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.image.io.polygon_loader.html b/cyto_dl.image.io.polygon_loader.html
index 3e66b9bf9..928ec785f 100644
--- a/cyto_dl.image.io.polygon_loader.html
+++ b/cyto_dl.image.io.polygon_loader.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.image.io.skimage_reader.html b/cyto_dl.image.io.skimage_reader.html
index f5eef9e32..f7869f36f 100644
--- a/cyto_dl.image.io.skimage_reader.html
+++ b/cyto_dl.image.io.skimage_reader.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.image.transforms.add_meta.html b/cyto_dl.image.transforms.add_meta.html
index 59a5c83cc..1cc13e1cb 100644
--- a/cyto_dl.image.transforms.add_meta.html
+++ b/cyto_dl.image.transforms.add_meta.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.image.transforms.bright_sampler.html b/cyto_dl.image.transforms.bright_sampler.html
index 9953578f4..1cf516975 100644
--- a/cyto_dl.image.transforms.bright_sampler.html
+++ b/cyto_dl.image.transforms.bright_sampler.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.image.transforms.clip.html b/cyto_dl.image.transforms.clip.html
index 192b606da..cfaf0e1fe 100644
--- a/cyto_dl.image.transforms.clip.html
+++ b/cyto_dl.image.transforms.clip.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.image.transforms.contrastadjust.html b/cyto_dl.image.transforms.contrastadjust.html
index 7fcb71f51..c40ddd27e 100644
--- a/cyto_dl.image.transforms.contrastadjust.html
+++ b/cyto_dl.image.transforms.contrastadjust.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.image.transforms.generate_jepa_masks.html b/cyto_dl.image.transforms.generate_jepa_masks.html
index 9bee639a4..9aef5f12e 100644
--- a/cyto_dl.image.transforms.generate_jepa_masks.html
+++ b/cyto_dl.image.transforms.generate_jepa_masks.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
@@ -438,7 +445,7 @@
 <span id="cyto-dl-image-transforms-generate-jepa-masks-module"></span><h1>cyto_dl.image.transforms.generate_jepa_masks module<a class="headerlink" href="#module-cyto_dl.image.transforms.generate_jepa_masks" title="Permalink to this heading">#</a></h1>
 <dl class="py class">
 <dt class="sig sig-object py" id="cyto_dl.image.transforms.generate_jepa_masks.JEPAMaskGenerator">
-<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">cyto_dl.image.transforms.generate_jepa_masks.</span></span><span class="sig-name descname"><span class="pre">JEPAMaskGenerator</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">mask_size</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">12</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">block_aspect_ratio</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">Tuple</span><span class="p"><span class="pre">[</span></span><span class="pre">float</span><span class="p"><span class="pre">]</span></span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">(0.5,</span> <span class="pre">1.5)</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">num_patches</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">Tuple</span><span class="p"><span class="pre">[</span></span><span class="pre">float</span><span class="p"><span class="pre">]</span></span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">(6,</span> <span class="pre">24,</span> <span class="pre">24)</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">mask_ratio</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">float</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">0.9</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/image/transforms/generate_jepa_masks.html#JEPAMaskGenerator"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.image.transforms.generate_jepa_masks.JEPAMaskGenerator" title="Permalink to this definition">#</a></dt>
+<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">cyto_dl.image.transforms.generate_jepa_masks.</span></span><span class="sig-name descname"><span class="pre">JEPAMaskGenerator</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">spatial_dims</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">mask_size</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">12</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">block_aspect_ratio</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">Tuple</span><span class="p"><span class="pre">[</span></span><span class="pre">float</span><span class="p"><span class="pre">]</span></span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">(0.5,</span> <span class="pre">1.5)</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">num_patches</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">Tuple</span><span class="p"><span class="pre">[</span></span><span class="pre">float</span><span class="p"><span class="pre">]</span></span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">(6,</span> <span class="pre">24,</span> <span class="pre">24)</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">mask_ratio</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">float</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">0.9</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/image/transforms/generate_jepa_masks.html#JEPAMaskGenerator"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.image.transforms.generate_jepa_masks.JEPAMaskGenerator" title="Permalink to this definition">#</a></dt>
 <dd><p>Bases: <code class="xref py py-class docutils literal notranslate"><span class="pre">RandomizableTransform</span></code></p>
 <p>Transform for generating Block-contiguous masks for JEPA training.</p>
 <p>This class works by randomly adding mask blocks until the mask_ratio is met or exceeded, then
@@ -446,6 +453,7 @@
 <dl class="field-list simple">
 <dt class="field-odd">Parameters<span class="colon">:</span></dt>
 <dd class="field-odd"><ul class="simple">
+<li><p><strong>spatial_dims</strong> (<em>int</em>) – The number of spatial dimensions of the image (2 or 3)</p></li>
 <li><p><strong>mask_size</strong> (<em>int, optional</em>) – The size of the blocks used to generate mask. Block dimensions are determined by the mask size and an aspect ratio sampled from the range  <cite>block_aspect_ratio</cite></p></li>
 <li><p><strong>block_aspect_ratio</strong> (<em>Tuple[float], optional</em>) – The low and high values for aspect ratio of the mask blocks</p></li>
 <li><p><strong>num_patches</strong> (<em>Tuple[int], optional</em>) – The number of patches used by the encoder for each dimension of the image (ZYX for 3D, YX for 2D)</p></li>
diff --git a/cyto_dl.image.transforms.html b/cyto_dl.image.transforms.html
index 97ad704e7..55c339b3c 100644
--- a/cyto_dl.image.transforms.html
+++ b/cyto_dl.image.transforms.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.image.transforms.merge.html b/cyto_dl.image.transforms.merge.html
index 02de85256..eb700be1e 100644
--- a/cyto_dl.image.transforms.merge.html
+++ b/cyto_dl.image.transforms.merge.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.image.transforms.multiscale_cropper.html b/cyto_dl.image.transforms.multiscale_cropper.html
index 6ecd0037b..517f10032 100644
--- a/cyto_dl.image.transforms.multiscale_cropper.html
+++ b/cyto_dl.image.transforms.multiscale_cropper.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.image.transforms.pad.html b/cyto_dl.image.transforms.pad.html
index d4095b165..498c55997 100644
--- a/cyto_dl.image.transforms.pad.html
+++ b/cyto_dl.image.transforms.pad.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.image.transforms.project.html b/cyto_dl.image.transforms.project.html
index 2f6921c86..84be4fa12 100644
--- a/cyto_dl.image.transforms.project.html
+++ b/cyto_dl.image.transforms.project.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.image.transforms.rotation_mask_transform.html b/cyto_dl.image.transforms.rotation_mask_transform.html
index d82a63351..dcf4f45b0 100644
--- a/cyto_dl.image.transforms.rotation_mask_transform.html
+++ b/cyto_dl.image.transforms.rotation_mask_transform.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.image.transforms.save.html b/cyto_dl.image.transforms.save.html
index 5af2db8a8..ff82bffb1 100644
--- a/cyto_dl.image.transforms.save.html
+++ b/cyto_dl.image.transforms.save.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.image.transforms.so2_random_rotation.html b/cyto_dl.image.transforms.so2_random_rotation.html
index 3173b563d..75c782c9a 100644
--- a/cyto_dl.image.transforms.so2_random_rotation.html
+++ b/cyto_dl.image.transforms.so2_random_rotation.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.image.transforms.track_transforms.html b/cyto_dl.image.transforms.track_transforms.html
index fc3fa5ede..d45591655 100644
--- a/cyto_dl.image.transforms.track_transforms.html
+++ b/cyto_dl.image.transforms.track_transforms.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.loggers.html b/cyto_dl.loggers.html
index d47f2c633..2c737f7a2 100644
--- a/cyto_dl.loggers.html
+++ b/cyto_dl.loggers.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.loggers.mlflow.html b/cyto_dl.loggers.mlflow.html
index 2e3ef2f40..6ed0c0df4 100644
--- a/cyto_dl.loggers.mlflow.html
+++ b/cyto_dl.loggers.mlflow.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.models.base_model.html b/cyto_dl.models.base_model.html
index 844391914..45395bbcd 100644
--- a/cyto_dl.models.base_model.html
+++ b/cyto_dl.models.base_model.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.models.basic_model.html b/cyto_dl.models.basic_model.html
index 13f94c034..f5a9775ce 100644
--- a/cyto_dl.models.basic_model.html
+++ b/cyto_dl.models.basic_model.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.models.classification.classification.html b/cyto_dl.models.classification.classification.html
index 3ba3699ef..0db308760 100644
--- a/cyto_dl.models.classification.classification.html
+++ b/cyto_dl.models.classification.classification.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.models.classification.html b/cyto_dl.models.classification.html
index a11bbb038..03503134e 100644
--- a/cyto_dl.models.classification.html
+++ b/cyto_dl.models.classification.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.models.classification.timepoint_classification.html b/cyto_dl.models.classification.timepoint_classification.html
index e546e57c8..88b546da7 100644
--- a/cyto_dl.models.classification.timepoint_classification.html
+++ b/cyto_dl.models.classification.timepoint_classification.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.models.contrastive.contrastive.html b/cyto_dl.models.contrastive.contrastive.html
index eff259463..a17c18578 100644
--- a/cyto_dl.models.contrastive.contrastive.html
+++ b/cyto_dl.models.contrastive.contrastive.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.models.contrastive.html b/cyto_dl.models.contrastive.html
index 28ec32bbb..af1c36562 100644
--- a/cyto_dl.models.contrastive.html
+++ b/cyto_dl.models.contrastive.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.models.handlers.base_handler.html b/cyto_dl.models.handlers.base_handler.html
index 7bc588bd6..39f0c7937 100644
--- a/cyto_dl.models.handlers.base_handler.html
+++ b/cyto_dl.models.handlers.base_handler.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.models.handlers.html b/cyto_dl.models.handlers.html
index e4ea0c43d..39290cc61 100644
--- a/cyto_dl.models.handlers.html
+++ b/cyto_dl.models.handlers.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.models.handlers.image_handler.html b/cyto_dl.models.handlers.image_handler.html
index cf17999e6..9f3180398 100644
--- a/cyto_dl.models.handlers.image_handler.html
+++ b/cyto_dl.models.handlers.image_handler.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.models.handlers.load_image_patch.html b/cyto_dl.models.handlers.load_image_patch.html
index 4b1bdd4c6..a4c7d12ad 100644
--- a/cyto_dl.models.handlers.load_image_patch.html
+++ b/cyto_dl.models.handlers.load_image_patch.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.models.html b/cyto_dl.models.html
index 884294ef6..b32433932 100644
--- a/cyto_dl.models.html
+++ b/cyto_dl.models.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.models.im2im.gan.html b/cyto_dl.models.im2im.gan.html
index 63173e0f9..f2f7fe64d 100644
--- a/cyto_dl.models.im2im.gan.html
+++ b/cyto_dl.models.im2im.gan.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.models.im2im.html b/cyto_dl.models.im2im.html
index 8701f6e72..c9f7a2c1f 100644
--- a/cyto_dl.models.im2im.html
+++ b/cyto_dl.models.im2im.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.models.im2im.multi_task.html b/cyto_dl.models.im2im.multi_task.html
index 92f1294f7..3b158b600 100644
--- a/cyto_dl.models.im2im.multi_task.html
+++ b/cyto_dl.models.im2im.multi_task.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.models.im2im.utils.html b/cyto_dl.models.im2im.utils.html
index fd7091303..17efb6dff 100644
--- a/cyto_dl.models.im2im.utils.html
+++ b/cyto_dl.models.im2im.utils.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.models.im2im.utils.instance_seg.html b/cyto_dl.models.im2im.utils.instance_seg.html
index 8809e6fb1..2252e5a6e 100644
--- a/cyto_dl.models.im2im.utils.instance_seg.html
+++ b/cyto_dl.models.im2im.utils.instance_seg.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.models.im2im.utils.noise_annealer.html b/cyto_dl.models.im2im.utils.noise_annealer.html
index 933f2d6a1..dcf5a26b2 100644
--- a/cyto_dl.models.im2im.utils.noise_annealer.html
+++ b/cyto_dl.models.im2im.utils.noise_annealer.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.models.im2im.utils.postprocessing.act_thresh_label.html b/cyto_dl.models.im2im.utils.postprocessing.act_thresh_label.html
index d93e72c13..9c99d1c82 100644
--- a/cyto_dl.models.im2im.utils.postprocessing.act_thresh_label.html
+++ b/cyto_dl.models.im2im.utils.postprocessing.act_thresh_label.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.models.im2im.utils.postprocessing.auto_thresh.html b/cyto_dl.models.im2im.utils.postprocessing.auto_thresh.html
index e4319d140..9b17ca71c 100644
--- a/cyto_dl.models.im2im.utils.postprocessing.auto_thresh.html
+++ b/cyto_dl.models.im2im.utils.postprocessing.auto_thresh.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.models.im2im.utils.postprocessing.dict_to_im.html b/cyto_dl.models.im2im.utils.postprocessing.dict_to_im.html
index 2ec5e16e7..ba0d355de 100644
--- a/cyto_dl.models.im2im.utils.postprocessing.dict_to_im.html
+++ b/cyto_dl.models.im2im.utils.postprocessing.dict_to_im.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.models.im2im.utils.postprocessing.html b/cyto_dl.models.im2im.utils.postprocessing.html
index 6e66ced0c..59208f818 100644
--- a/cyto_dl.models.im2im.utils.postprocessing.html
+++ b/cyto_dl.models.im2im.utils.postprocessing.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.models.jepa.html b/cyto_dl.models.jepa.html
index 78851711d..64f284970 100644
--- a/cyto_dl.models.jepa.html
+++ b/cyto_dl.models.jepa.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.models.jepa.ijepa.html b/cyto_dl.models.jepa.ijepa.html
index c19d62e94..a1781b205 100644
--- a/cyto_dl.models.jepa.ijepa.html
+++ b/cyto_dl.models.jepa.ijepa.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.models.jepa.iwm.html b/cyto_dl.models.jepa.iwm.html
index cee324fcb..83bd80c61 100644
--- a/cyto_dl.models.jepa.iwm.html
+++ b/cyto_dl.models.jepa.iwm.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.models.jepa.jepa_base.html b/cyto_dl.models.jepa.jepa_base.html
index a20302032..3d201649b 100644
--- a/cyto_dl.models.jepa.jepa_base.html
+++ b/cyto_dl.models.jepa.jepa_base.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.models.utils.html b/cyto_dl.models.utils.html
index ece30d129..589cf49b0 100644
--- a/cyto_dl.models.utils.html
+++ b/cyto_dl.models.utils.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.models.utils.mlflow.html b/cyto_dl.models.utils.mlflow.html
index c546f2066..184518bf5 100644
--- a/cyto_dl.models.utils.mlflow.html
+++ b/cyto_dl.models.utils.mlflow.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.models.utils.utils.html b/cyto_dl.models.utils.utils.html
index 4af8c2956..96f33cf6e 100644
--- a/cyto_dl.models.utils.utils.html
+++ b/cyto_dl.models.utils.utils.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.models.vae.base_vae.html b/cyto_dl.models.vae.base_vae.html
index e03180070..d46e88477 100644
--- a/cyto_dl.models.vae.base_vae.html
+++ b/cyto_dl.models.vae.base_vae.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.models.vae.html b/cyto_dl.models.vae.html
index 617546708..84fa86719 100644
--- a/cyto_dl.models.vae.html
+++ b/cyto_dl.models.vae.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.models.vae.image_canon_vae.html b/cyto_dl.models.vae.image_canon_vae.html
index f25968271..8bcf5f6de 100644
--- a/cyto_dl.models.vae.image_canon_vae.html
+++ b/cyto_dl.models.vae.image_canon_vae.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.models.vae.image_encoder.html b/cyto_dl.models.vae.image_encoder.html
index 29f2793dd..0bddb00a0 100644
--- a/cyto_dl.models.vae.image_encoder.html
+++ b/cyto_dl.models.vae.image_encoder.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.models.vae.image_vae.html b/cyto_dl.models.vae.image_vae.html
index a3ede2091..4c5d0a130 100644
--- a/cyto_dl.models.vae.image_vae.html
+++ b/cyto_dl.models.vae.image_vae.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.models.vae.latent_loss_vae.html b/cyto_dl.models.vae.latent_loss_vae.html
index d8a5a735d..5af211da0 100644
--- a/cyto_dl.models.vae.latent_loss_vae.html
+++ b/cyto_dl.models.vae.latent_loss_vae.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.models.vae.o2_spharm_vae.html b/cyto_dl.models.vae.o2_spharm_vae.html
index 830ffe088..aa7c734d4 100644
--- a/cyto_dl.models.vae.o2_spharm_vae.html
+++ b/cyto_dl.models.vae.o2_spharm_vae.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.models.vae.o2_spharm_vae.o2_spharm_encoder.html b/cyto_dl.models.vae.o2_spharm_vae.o2_spharm_encoder.html
index 4e588f820..fe95f5b15 100644
--- a/cyto_dl.models.vae.o2_spharm_vae.o2_spharm_encoder.html
+++ b/cyto_dl.models.vae.o2_spharm_vae.o2_spharm_encoder.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.models.vae.o2_spharm_vae.o2_spharm_vae.html b/cyto_dl.models.vae.o2_spharm_vae.o2_spharm_vae.html
index f397d2a5f..e13159088 100644
--- a/cyto_dl.models.vae.o2_spharm_vae.o2_spharm_vae.html
+++ b/cyto_dl.models.vae.o2_spharm_vae.o2_spharm_vae.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.models.vae.point_cloud_vae.html b/cyto_dl.models.vae.point_cloud_vae.html
index ae1abe738..a172df2ce 100644
--- a/cyto_dl.models.vae.point_cloud_vae.html
+++ b/cyto_dl.models.vae.point_cloud_vae.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.models.vae.priors.abstract_prior.html b/cyto_dl.models.vae.priors.abstract_prior.html
index ce3811dcd..4ae142020 100644
--- a/cyto_dl.models.vae.priors.abstract_prior.html
+++ b/cyto_dl.models.vae.priors.abstract_prior.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.models.vae.priors.gaussian.html b/cyto_dl.models.vae.priors.gaussian.html
index 63eb71eea..265acef28 100644
--- a/cyto_dl.models.vae.priors.gaussian.html
+++ b/cyto_dl.models.vae.priors.gaussian.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.models.vae.priors.html b/cyto_dl.models.vae.priors.html
index b043f42bf..9e72e61e1 100644
--- a/cyto_dl.models.vae.priors.html
+++ b/cyto_dl.models.vae.priors.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.models.vae.priors.identity_prior.html b/cyto_dl.models.vae.priors.identity_prior.html
index ff846286f..37282883d 100644
--- a/cyto_dl.models.vae.priors.identity_prior.html
+++ b/cyto_dl.models.vae.priors.identity_prior.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.models.vae.priors.joint_prior.html b/cyto_dl.models.vae.priors.joint_prior.html
index 12a29a538..b7fb087c2 100644
--- a/cyto_dl.models.vae.priors.joint_prior.html
+++ b/cyto_dl.models.vae.priors.joint_prior.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.models.vae.tabular_vae.html b/cyto_dl.models.vae.tabular_vae.html
index 49024f542..f40edf3df 100644
--- a/cyto_dl.models.vae.tabular_vae.html
+++ b/cyto_dl.models.vae.tabular_vae.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.nn.discriminators.html b/cyto_dl.nn.discriminators.html
index 2d29b1750..410994539 100644
--- a/cyto_dl.nn.discriminators.html
+++ b/cyto_dl.nn.discriminators.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.nn.discriminators.multi_scale_discriminator.html b/cyto_dl.nn.discriminators.multi_scale_discriminator.html
index 3f02c0344..0005257f8 100644
--- a/cyto_dl.nn.discriminators.multi_scale_discriminator.html
+++ b/cyto_dl.nn.discriminators.multi_scale_discriminator.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.nn.discriminators.n_layer_discriminator.html b/cyto_dl.nn.discriminators.n_layer_discriminator.html
index 2ae9e8d0b..731711445 100644
--- a/cyto_dl.nn.discriminators.n_layer_discriminator.html
+++ b/cyto_dl.nn.discriminators.n_layer_discriminator.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.nn.head.base_head.html b/cyto_dl.nn.head.base_head.html
index c8d9cee89..1ddf905c6 100644
--- a/cyto_dl.nn.head.base_head.html
+++ b/cyto_dl.nn.head.base_head.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.nn.head.gan_head.html b/cyto_dl.nn.head.gan_head.html
index 0568f020d..bb473f010 100644
--- a/cyto_dl.nn.head.gan_head.html
+++ b/cyto_dl.nn.head.gan_head.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.nn.head.gan_head_superres.html b/cyto_dl.nn.head.gan_head_superres.html
index b0eaff9fd..e50e84353 100644
--- a/cyto_dl.nn.head.gan_head_superres.html
+++ b/cyto_dl.nn.head.gan_head_superres.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.nn.head.html b/cyto_dl.nn.head.html
index 5c440cc9b..c87eaaf69 100644
--- a/cyto_dl.nn.head.html
+++ b/cyto_dl.nn.head.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.nn.head.mae_head.html b/cyto_dl.nn.head.mae_head.html
index 759cebd24..b706cb450 100644
--- a/cyto_dl.nn.head.mae_head.html
+++ b/cyto_dl.nn.head.mae_head.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.nn.head.mask_head.html b/cyto_dl.nn.head.mask_head.html
index 1f9c000df..254512c17 100644
--- a/cyto_dl.nn.head.mask_head.html
+++ b/cyto_dl.nn.head.mask_head.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.nn.head.res_blocks_head.html b/cyto_dl.nn.head.res_blocks_head.html
index 8e0b48a10..7aef7d9cb 100644
--- a/cyto_dl.nn.head.res_blocks_head.html
+++ b/cyto_dl.nn.head.res_blocks_head.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.nn.head.vic_reg.html b/cyto_dl.nn.head.vic_reg.html
index 8d95fe422..e685d6e6b 100644
--- a/cyto_dl.nn.head.vic_reg.html
+++ b/cyto_dl.nn.head.vic_reg.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.nn.hr_skip.html b/cyto_dl.nn.hr_skip.html
index ff1e550ae..6e1ac2310 100644
--- a/cyto_dl.nn.hr_skip.html
+++ b/cyto_dl.nn.hr_skip.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.nn.html b/cyto_dl.nn.html
index 43c213c9c..07531e7ae 100644
--- a/cyto_dl.nn.html
+++ b/cyto_dl.nn.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
@@ -557,26 +564,34 @@ <h2>Subpackages<a class="headerlink" href="#subpackages" title="Permalink to thi
 <li class="toctree-l1"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><ul>
 <li class="toctree-l2"><a class="reference internal" href="cyto_dl.nn.vits.html#subpackages">Subpackages</a><ul>
 <li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><ul>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html#subpackages">Subpackages</a></li>
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html#submodules">Submodules</a></li>
 </ul>
 </li>
 </ul>
 </li>
 <li class="toctree-l2"><a class="reference internal" href="cyto_dl.nn.vits.html#submodules">Submodules</a><ul>
-<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a><ul>
-<li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html#cyto_dl.nn.vits.cross_mae.CrossMAE_Decoder"><code class="docutils literal notranslate"><span class="pre">CrossMAE_Decoder</span></code></a></li>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a><ul>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html#cyto_dl.nn.vits.decoder.CrossMAE_Decoder"><code class="docutils literal notranslate"><span class="pre">CrossMAE_Decoder</span></code></a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html#cyto_dl.nn.vits.decoder.MAE_Decoder"><code class="docutils literal notranslate"><span class="pre">MAE_Decoder</span></code></a></li>
 </ul>
 </li>
-<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a><ul>
-<li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html#cyto_dl.nn.vits.jepa.IWMPredictor"><code class="docutils literal notranslate"><span class="pre">IWMPredictor</span></code></a></li>
-<li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html#cyto_dl.nn.vits.jepa.JEPAEncoder"><code class="docutils literal notranslate"><span class="pre">JEPAEncoder</span></code></a></li>
-<li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html#cyto_dl.nn.vits.jepa.JEPAPredictor"><code class="docutils literal notranslate"><span class="pre">JEPAPredictor</span></code></a></li>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a><ul>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html#cyto_dl.nn.vits.encoder.HieraEncoder"><code class="docutils literal notranslate"><span class="pre">HieraEncoder</span></code></a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html#cyto_dl.nn.vits.encoder.JEPAEncoder"><code class="docutils literal notranslate"><span class="pre">JEPAEncoder</span></code></a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html#cyto_dl.nn.vits.encoder.MAE_Encoder"><code class="docutils literal notranslate"><span class="pre">MAE_Encoder</span></code></a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html#cyto_dl.nn.vits.encoder.SpatialMerger"><code class="docutils literal notranslate"><span class="pre">SpatialMerger</span></code></a></li>
 </ul>
 </li>
 <li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a><ul>
-<li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.MAE_Decoder"><code class="docutils literal notranslate"><span class="pre">MAE_Decoder</span></code></a></li>
-<li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.MAE_Encoder"><code class="docutils literal notranslate"><span class="pre">MAE_Encoder</span></code></a></li>
-<li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.MAE_ViT"><code class="docutils literal notranslate"><span class="pre">MAE_ViT</span></code></a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.HieraMAE"><code class="docutils literal notranslate"><span class="pre">HieraMAE</span></code></a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.MAE"><code class="docutils literal notranslate"><span class="pre">MAE</span></code></a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.MAE_Base"><code class="docutils literal notranslate"><span class="pre">MAE_Base</span></code></a></li>
+</ul>
+</li>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a><ul>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html#cyto_dl.nn.vits.predictor.IWMPredictor"><code class="docutils literal notranslate"><span class="pre">IWMPredictor</span></code></a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html#cyto_dl.nn.vits.predictor.JEPAPredictor"><code class="docutils literal notranslate"><span class="pre">JEPAPredictor</span></code></a></li>
 </ul>
 </li>
 <li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a><ul>
@@ -587,6 +602,8 @@ <h2>Subpackages<a class="headerlink" href="#subpackages" title="Permalink to thi
 </li>
 <li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a><ul>
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.vits.utils.html#cyto_dl.nn.vits.utils.get_positional_embedding"><code class="docutils literal notranslate"><span class="pre">get_positional_embedding()</span></code></a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.vits.utils.html#cyto_dl.nn.vits.utils.match_tuple_dimensions"><code class="docutils literal notranslate"><span class="pre">match_tuple_dimensions()</span></code></a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.vits.utils.html#cyto_dl.nn.vits.utils.random_indexes"><code class="docutils literal notranslate"><span class="pre">random_indexes()</span></code></a></li>
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.vits.utils.html#cyto_dl.nn.vits.utils.take_indexes"><code class="docutils literal notranslate"><span class="pre">take_indexes()</span></code></a></li>
 </ul>
 </li>
diff --git a/cyto_dl.nn.losses.adversarial_loss.html b/cyto_dl.nn.losses.adversarial_loss.html
index 49013d883..a5f9816f3 100644
--- a/cyto_dl.nn.losses.adversarial_loss.html
+++ b/cyto_dl.nn.losses.adversarial_loss.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.nn.losses.chamfer_loss.html b/cyto_dl.nn.losses.chamfer_loss.html
index ade64ee0d..ef10b5185 100644
--- a/cyto_dl.nn.losses.chamfer_loss.html
+++ b/cyto_dl.nn.losses.chamfer_loss.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.nn.losses.continuous_bernoulli.html b/cyto_dl.nn.losses.continuous_bernoulli.html
index ae68705d9..86a820c34 100644
--- a/cyto_dl.nn.losses.continuous_bernoulli.html
+++ b/cyto_dl.nn.losses.continuous_bernoulli.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.nn.losses.cosine_loss.html b/cyto_dl.nn.losses.cosine_loss.html
index 8cd783220..961310b9a 100644
--- a/cyto_dl.nn.losses.cosine_loss.html
+++ b/cyto_dl.nn.losses.cosine_loss.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.nn.losses.gan_loss.html b/cyto_dl.nn.losses.gan_loss.html
index 88ea86e8c..8ba634da6 100644
--- a/cyto_dl.nn.losses.gan_loss.html
+++ b/cyto_dl.nn.losses.gan_loss.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.nn.losses.gaussian_nll_loss.html b/cyto_dl.nn.losses.gaussian_nll_loss.html
index dc95b1554..489eb4d5c 100644
--- a/cyto_dl.nn.losses.gaussian_nll_loss.html
+++ b/cyto_dl.nn.losses.gaussian_nll_loss.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.nn.losses.geomloss.html b/cyto_dl.nn.losses.geomloss.html
index 36b46585f..485aba1ee 100644
--- a/cyto_dl.nn.losses.geomloss.html
+++ b/cyto_dl.nn.losses.geomloss.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.nn.losses.html b/cyto_dl.nn.losses.html
index 136b3af3f..3216e8cf1 100644
--- a/cyto_dl.nn.losses.html
+++ b/cyto_dl.nn.losses.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.nn.losses.loss_wrapper.html b/cyto_dl.nn.losses.loss_wrapper.html
index a7d340fee..7180d0a7c 100644
--- a/cyto_dl.nn.losses.loss_wrapper.html
+++ b/cyto_dl.nn.losses.loss_wrapper.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.nn.losses.spharm_loss.html b/cyto_dl.nn.losses.spharm_loss.html
index f2bd17fa7..919e3822f 100644
--- a/cyto_dl.nn.losses.spharm_loss.html
+++ b/cyto_dl.nn.losses.spharm_loss.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.nn.losses.threshold_loss.html b/cyto_dl.nn.losses.threshold_loss.html
index 0f981d1fd..29b690aa7 100644
--- a/cyto_dl.nn.losses.threshold_loss.html
+++ b/cyto_dl.nn.losses.threshold_loss.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.nn.losses.vic_reg.html b/cyto_dl.nn.losses.vic_reg.html
index a7fe56a93..e75f1a1c1 100644
--- a/cyto_dl.nn.losses.vic_reg.html
+++ b/cyto_dl.nn.losses.vic_reg.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.nn.losses.weibull.html b/cyto_dl.nn.losses.weibull.html
index 09bf3fa02..e586b5b8d 100644
--- a/cyto_dl.nn.losses.weibull.html
+++ b/cyto_dl.nn.losses.weibull.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.nn.losses.weighted_mse_loss.html b/cyto_dl.nn.losses.weighted_mse_loss.html
index 45471eee1..ada7e53ff 100644
--- a/cyto_dl.nn.losses.weighted_mse_loss.html
+++ b/cyto_dl.nn.losses.weighted_mse_loss.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.nn.mlp.html b/cyto_dl.nn.mlp.html
index 937342519..ee276ba68 100644
--- a/cyto_dl.nn.mlp.html
+++ b/cyto_dl.nn.mlp.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.nn.point_cloud.dgcnn.html b/cyto_dl.nn.point_cloud.dgcnn.html
index 4c8427b61..34e203774 100644
--- a/cyto_dl.nn.point_cloud.dgcnn.html
+++ b/cyto_dl.nn.point_cloud.dgcnn.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.nn.point_cloud.folding_net.html b/cyto_dl.nn.point_cloud.folding_net.html
index 0f5939c9e..bc3a00a9f 100644
--- a/cyto_dl.nn.point_cloud.folding_net.html
+++ b/cyto_dl.nn.point_cloud.folding_net.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.nn.point_cloud.graph_functions.html b/cyto_dl.nn.point_cloud.graph_functions.html
index 1df6cf943..b8fa07ee3 100644
--- a/cyto_dl.nn.point_cloud.graph_functions.html
+++ b/cyto_dl.nn.point_cloud.graph_functions.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.nn.point_cloud.html b/cyto_dl.nn.point_cloud.html
index a91b4d0b3..37d1eff75 100644
--- a/cyto_dl.nn.point_cloud.html
+++ b/cyto_dl.nn.point_cloud.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.nn.point_cloud.vnn.html b/cyto_dl.nn.point_cloud.vnn.html
index 92c821723..6720b4764 100644
--- a/cyto_dl.nn.point_cloud.vnn.html
+++ b/cyto_dl.nn.point_cloud.vnn.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.nn.res_unit.html b/cyto_dl.nn.res_unit.html
index a665f4870..5a3ab7d77 100644
--- a/cyto_dl.nn.res_unit.html
+++ b/cyto_dl.nn.res_unit.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.nn.spatial_transformer.html b/cyto_dl.nn.spatial_transformer.html
index 7cb18b902..d4b6e6243 100644
--- a/cyto_dl.nn.spatial_transformer.html
+++ b/cyto_dl.nn.spatial_transformer.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.nn.track_sequence_predictor.html b/cyto_dl.nn.track_sequence_predictor.html
index 39349b00b..f65d8f809 100644
--- a/cyto_dl.nn.track_sequence_predictor.html
+++ b/cyto_dl.nn.track_sequence_predictor.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4 current current-page"><a class="current reference internal" href="#">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.nn.vits.blocks.cross_attention.html b/cyto_dl.nn.vits.blocks.cross_attention.html
index 7354c7a4e..11545474a 100644
--- a/cyto_dl.nn.vits.blocks.cross_attention.html
+++ b/cyto_dl.nn.vits.blocks.cross_attention.html
@@ -3,7 +3,7 @@
   <head><meta charset="utf-8"/>
     <meta name="viewport" content="width=device-width,initial-scale=1"/>
     <meta name="color-scheme" content="light dark"><meta name="generator" content="Docutils 0.19: https://docutils.sourceforge.io/" />
-<link rel="index" title="Index" href="genindex.html" /><link rel="search" title="Search" href="search.html" /><link rel="next" title="cyto_dl.nn.vits.blocks.intermediate_weigher module" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html" /><link rel="prev" title="cyto_dl.nn.vits.blocks package" href="cyto_dl.nn.vits.blocks.html" />
+<link rel="index" title="Index" href="genindex.html" /><link rel="search" title="Search" href="search.html" /><link rel="next" title="cyto_dl.nn.vits.blocks.intermediate_weigher module" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html" /><link rel="prev" title="cyto_dl.nn.vits.blocks.patchify.patchify_hiera module" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html" />
 
     <!-- Generated with Sphinx 6.1.3 and Furo 2023.03.23 -->
         <title>cyto_dl.nn.vits.blocks.cross_attention module - cyto-dl 0.4.0 documentation</title>
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 current has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
 <li class="toctree-l5 current has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6 current current-page"><a class="current reference internal" href="#">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
@@ -499,14 +506,14 @@
               </div>
               <svg class="furo-related-icon"><use href="#svg-arrow-right"></use></svg>
             </a>
-          <a class="prev-page" href="cyto_dl.nn.vits.blocks.html">
+          <a class="prev-page" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">
               <svg class="furo-related-icon"><use href="#svg-arrow-right"></use></svg>
               <div class="page-info">
                 <div class="context">
                   <span>Previous</span>
                 </div>
                 
-                <div class="title">cyto_dl.nn.vits.blocks package</div>
+                <div class="title">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</div>
                 
               </div>
             </a>
diff --git a/cyto_dl.nn.vits.blocks.html b/cyto_dl.nn.vits.blocks.html
index a1c62b4ef..9d7d84fe6 100644
--- a/cyto_dl.nn.vits.blocks.html
+++ b/cyto_dl.nn.vits.blocks.html
@@ -3,7 +3,7 @@
   <head><meta charset="utf-8"/>
     <meta name="viewport" content="width=device-width,initial-scale=1"/>
     <meta name="color-scheme" content="light dark"><meta name="generator" content="Docutils 0.19: https://docutils.sourceforge.io/" />
-<link rel="index" title="Index" href="genindex.html" /><link rel="search" title="Search" href="search.html" /><link rel="next" title="cyto_dl.nn.vits.blocks.cross_attention module" href="cyto_dl.nn.vits.blocks.cross_attention.html" /><link rel="prev" title="cyto_dl.nn.vits package" href="cyto_dl.nn.vits.html" />
+<link rel="index" title="Index" href="genindex.html" /><link rel="search" title="Search" href="search.html" /><link rel="next" title="cyto_dl.nn.vits.blocks.patchify package" href="cyto_dl.nn.vits.blocks.patchify.html" /><link rel="prev" title="cyto_dl.nn.vits package" href="cyto_dl.nn.vits.html" />
 
     <!-- Generated with Sphinx 6.1.3 and Furo 2023.03.23 -->
         <title>cyto_dl.nn.vits.blocks package - cyto-dl 0.4.0 documentation</title>
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 current has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
 <li class="toctree-l5 current has-children current-page"><a class="current reference internal" href="#">cyto_dl.nn.vits.blocks package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
@@ -436,6 +443,32 @@
         <article role="main">
           <section id="module-cyto_dl.nn.vits.blocks">
 <span id="cyto-dl-nn-vits-blocks-package"></span><h1>cyto_dl.nn.vits.blocks package<a class="headerlink" href="#module-cyto_dl.nn.vits.blocks" title="Permalink to this heading">#</a></h1>
+<section id="subpackages">
+<h2>Subpackages<a class="headerlink" href="#subpackages" title="Permalink to this heading">#</a></h2>
+<div class="toctree-wrapper compound">
+<ul>
+<li class="toctree-l1"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><ul>
+<li class="toctree-l2"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html#submodules">Submodules</a><ul>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a><ul>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html#cyto_dl.nn.vits.blocks.patchify.patchify.Patchify"><code class="docutils literal notranslate"><span class="pre">Patchify</span></code></a></li>
+</ul>
+</li>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a><ul>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html#cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase"><code class="docutils literal notranslate"><span class="pre">PatchifyBase</span></code></a></li>
+</ul>
+</li>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a><ul>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html#cyto_dl.nn.vits.blocks.patchify.patchify_hiera.PatchifyHiera"><code class="docutils literal notranslate"><span class="pre">PatchifyHiera</span></code></a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html#cyto_dl.nn.vits.blocks.patchify.patchify_hiera.take_indexes_mask"><code class="docutils literal notranslate"><span class="pre">take_indexes_mask()</span></code></a></li>
+</ul>
+</li>
+</ul>
+</li>
+</ul>
+</li>
+</ul>
+</div>
+</section>
 <section id="submodules">
 <h2>Submodules<a class="headerlink" href="#submodules" title="Permalink to this heading">#</a></h2>
 <div class="toctree-wrapper compound">
@@ -466,13 +499,15 @@ <h2>Submodules<a class="headerlink" href="#submodules" title="Permalink to this
 </li>
 </ul>
 </li>
-<li class="toctree-l1"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a><ul>
-<li class="toctree-l2"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html#cyto_dl.nn.vits.blocks.patchify.Patchify"><code class="docutils literal notranslate"><span class="pre">Patchify</span></code></a><ul>
-<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html#cyto_dl.nn.vits.blocks.patchify.Patchify.forward"><code class="docutils literal notranslate"><span class="pre">Patchify.forward()</span></code></a></li>
-<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html#cyto_dl.nn.vits.blocks.patchify.Patchify.get_mask"><code class="docutils literal notranslate"><span class="pre">Patchify.get_mask()</span></code></a></li>
+<li class="toctree-l1"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a><ul>
+<li class="toctree-l2"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html#cyto_dl.nn.vits.blocks.masked_unit_attention.HieraBlock"><code class="docutils literal notranslate"><span class="pre">HieraBlock</span></code></a><ul>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html#cyto_dl.nn.vits.blocks.masked_unit_attention.HieraBlock.forward"><code class="docutils literal notranslate"><span class="pre">HieraBlock.forward()</span></code></a></li>
+</ul>
+</li>
+<li class="toctree-l2"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html#cyto_dl.nn.vits.blocks.masked_unit_attention.MaskUnitAttention"><code class="docutils literal notranslate"><span class="pre">MaskUnitAttention</span></code></a><ul>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html#cyto_dl.nn.vits.blocks.masked_unit_attention.MaskUnitAttention.forward"><code class="docutils literal notranslate"><span class="pre">MaskUnitAttention.forward()</span></code></a></li>
 </ul>
 </li>
-<li class="toctree-l2"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html#cyto_dl.nn.vits.blocks.patchify.random_indexes"><code class="docutils literal notranslate"><span class="pre">random_indexes()</span></code></a></li>
 </ul>
 </li>
 </ul>
@@ -485,12 +520,12 @@ <h2>Submodules<a class="headerlink" href="#submodules" title="Permalink to this
       <footer>
         
         <div class="related-pages">
-          <a class="next-page" href="cyto_dl.nn.vits.blocks.cross_attention.html">
+          <a class="next-page" href="cyto_dl.nn.vits.blocks.patchify.html">
               <div class="page-info">
                 <div class="context">
                   <span>Next</span>
                 </div>
-                <div class="title">cyto_dl.nn.vits.blocks.cross_attention module</div>
+                <div class="title">cyto_dl.nn.vits.blocks.patchify package</div>
               </div>
               <svg class="furo-related-icon"><use href="#svg-arrow-right"></use></svg>
             </a>
@@ -536,6 +571,7 @@ <h2>Submodules<a class="headerlink" href="#submodules" title="Permalink to this
           <div class="toc-tree">
             <ul>
 <li><a class="reference internal" href="#">cyto_dl.nn.vits.blocks package</a><ul>
+<li><a class="reference internal" href="#subpackages">Subpackages</a></li>
 <li><a class="reference internal" href="#submodules">Submodules</a></li>
 </ul>
 </li>
diff --git a/cyto_dl.nn.vits.blocks.intermediate_weigher.html b/cyto_dl.nn.vits.blocks.intermediate_weigher.html
index 748c4dc32..7c90bc8f2 100644
--- a/cyto_dl.nn.vits.blocks.intermediate_weigher.html
+++ b/cyto_dl.nn.vits.blocks.intermediate_weigher.html
@@ -3,7 +3,7 @@
   <head><meta charset="utf-8"/>
     <meta name="viewport" content="width=device-width,initial-scale=1"/>
     <meta name="color-scheme" content="light dark"><meta name="generator" content="Docutils 0.19: https://docutils.sourceforge.io/" />
-<link rel="index" title="Index" href="genindex.html" /><link rel="search" title="Search" href="search.html" /><link rel="next" title="cyto_dl.nn.vits.blocks.patchify module" href="cyto_dl.nn.vits.blocks.patchify.html" /><link rel="prev" title="cyto_dl.nn.vits.blocks.cross_attention module" href="cyto_dl.nn.vits.blocks.cross_attention.html" />
+<link rel="index" title="Index" href="genindex.html" /><link rel="search" title="Search" href="search.html" /><link rel="next" title="cyto_dl.nn.vits.blocks.masked_unit_attention module" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html" /><link rel="prev" title="cyto_dl.nn.vits.blocks.cross_attention module" href="cyto_dl.nn.vits.blocks.cross_attention.html" />
 
     <!-- Generated with Sphinx 6.1.3 and Furo 2023.03.23 -->
         <title>cyto_dl.nn.vits.blocks.intermediate_weigher module - cyto-dl 0.4.0 documentation</title>
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 current has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
 <li class="toctree-l5 current has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6 current current-page"><a class="current reference internal" href="#">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
@@ -456,12 +463,12 @@
       <footer>
         
         <div class="related-pages">
-          <a class="next-page" href="cyto_dl.nn.vits.blocks.patchify.html">
+          <a class="next-page" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">
               <div class="page-info">
                 <div class="context">
                   <span>Next</span>
                 </div>
-                <div class="title">cyto_dl.nn.vits.blocks.patchify module</div>
+                <div class="title">cyto_dl.nn.vits.blocks.masked_unit_attention module</div>
               </div>
               <svg class="furo-related-icon"><use href="#svg-arrow-right"></use></svg>
             </a>
diff --git a/cyto_dl.nn.vits.blocks.masked_unit_attention.html b/cyto_dl.nn.vits.blocks.masked_unit_attention.html
new file mode 100644
index 000000000..4653f69a7
--- /dev/null
+++ b/cyto_dl.nn.vits.blocks.masked_unit_attention.html
@@ -0,0 +1,582 @@
+<!doctype html>
+<html class="no-js" lang="en">
+  <head><meta charset="utf-8"/>
+    <meta name="viewport" content="width=device-width,initial-scale=1"/>
+    <meta name="color-scheme" content="light dark"><meta name="generator" content="Docutils 0.19: https://docutils.sourceforge.io/" />
+<link rel="index" title="Index" href="genindex.html" /><link rel="search" title="Search" href="search.html" /><link rel="next" title="cyto_dl.nn.vits.decoder module" href="cyto_dl.nn.vits.decoder.html" /><link rel="prev" title="cyto_dl.nn.vits.blocks.intermediate_weigher module" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html" />
+
+    <!-- Generated with Sphinx 6.1.3 and Furo 2023.03.23 -->
+        <title>cyto_dl.nn.vits.blocks.masked_unit_attention module - cyto-dl 0.4.0 documentation</title>
+      <link rel="stylesheet" type="text/css" href="_static/pygments.css" />
+    <link rel="stylesheet" type="text/css" href="_static/styles/furo.css?digest=fad236701ea90a88636c2a8c73b44ae642ed2a53" />
+    <link rel="stylesheet" type="text/css" href="_static/styles/furo-extensions.css?digest=30d1aed668e5c3a91c3e3bf6a60b675221979f0e" />
+    
+    
+
+
+<style>
+  body {
+    --color-code-background: #eeffcc;
+  --color-code-foreground: black;
+  
+  }
+  @media not print {
+    body[data-theme="dark"] {
+      --color-code-background: #202020;
+  --color-code-foreground: #d0d0d0;
+  
+    }
+    @media (prefers-color-scheme: dark) {
+      body:not([data-theme="light"]) {
+        --color-code-background: #202020;
+  --color-code-foreground: #d0d0d0;
+  
+      }
+    }
+  }
+</style></head>
+  <body>
+    
+    <script>
+      document.body.dataset.theme = localStorage.getItem("theme") || "auto";
+    </script>
+    
+
+<svg xmlns="http://www.w3.org/2000/svg" style="display: none;">
+  <symbol id="svg-toc" viewBox="0 0 24 24">
+    <title>Contents</title>
+    <svg stroke="currentColor" fill="currentColor" stroke-width="0" viewBox="0 0 1024 1024">
+      <path d="M408 442h480c4.4 0 8-3.6 8-8v-56c0-4.4-3.6-8-8-8H408c-4.4 0-8 3.6-8 8v56c0 4.4 3.6 8 8 8zm-8 204c0 4.4 3.6 8 8 8h480c4.4 0 8-3.6 8-8v-56c0-4.4-3.6-8-8-8H408c-4.4 0-8 3.6-8 8v56zm504-486H120c-4.4 0-8 3.6-8 8v56c0 4.4 3.6 8 8 8h784c4.4 0 8-3.6 8-8v-56c0-4.4-3.6-8-8-8zm0 632H120c-4.4 0-8 3.6-8 8v56c0 4.4 3.6 8 8 8h784c4.4 0 8-3.6 8-8v-56c0-4.4-3.6-8-8-8zM115.4 518.9L271.7 642c5.8 4.6 14.4.5 14.4-6.9V388.9c0-7.4-8.5-11.5-14.4-6.9L115.4 505.1a8.74 8.74 0 0 0 0 13.8z"/>
+    </svg>
+  </symbol>
+  <symbol id="svg-menu" viewBox="0 0 24 24">
+    <title>Menu</title>
+    <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
+      stroke-width="2" stroke-linecap="round" stroke-linejoin="round" class="feather-menu">
+      <line x1="3" y1="12" x2="21" y2="12"></line>
+      <line x1="3" y1="6" x2="21" y2="6"></line>
+      <line x1="3" y1="18" x2="21" y2="18"></line>
+    </svg>
+  </symbol>
+  <symbol id="svg-arrow-right" viewBox="0 0 24 24">
+    <title>Expand</title>
+    <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
+      stroke-width="2" stroke-linecap="round" stroke-linejoin="round" class="feather-chevron-right">
+      <polyline points="9 18 15 12 9 6"></polyline>
+    </svg>
+  </symbol>
+  <symbol id="svg-sun" viewBox="0 0 24 24">
+    <title>Light mode</title>
+    <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
+      stroke-width="1.5" stroke-linecap="round" stroke-linejoin="round" class="feather-sun">
+      <circle cx="12" cy="12" r="5"></circle>
+      <line x1="12" y1="1" x2="12" y2="3"></line>
+      <line x1="12" y1="21" x2="12" y2="23"></line>
+      <line x1="4.22" y1="4.22" x2="5.64" y2="5.64"></line>
+      <line x1="18.36" y1="18.36" x2="19.78" y2="19.78"></line>
+      <line x1="1" y1="12" x2="3" y2="12"></line>
+      <line x1="21" y1="12" x2="23" y2="12"></line>
+      <line x1="4.22" y1="19.78" x2="5.64" y2="18.36"></line>
+      <line x1="18.36" y1="5.64" x2="19.78" y2="4.22"></line>
+    </svg>
+  </symbol>
+  <symbol id="svg-moon" viewBox="0 0 24 24">
+    <title>Dark mode</title>
+    <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
+      stroke-width="1.5" stroke-linecap="round" stroke-linejoin="round" class="icon-tabler-moon">
+      <path stroke="none" d="M0 0h24v24H0z" fill="none" />
+      <path d="M12 3c.132 0 .263 0 .393 0a7.5 7.5 0 0 0 7.92 12.446a9 9 0 1 1 -8.313 -12.454z" />
+    </svg>
+  </symbol>
+  <symbol id="svg-sun-half" viewBox="0 0 24 24">
+    <title>Auto light/dark mode</title>
+    <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
+      stroke-width="1.5" stroke-linecap="round" stroke-linejoin="round" class="icon-tabler-shadow">
+      <path stroke="none" d="M0 0h24v24H0z" fill="none"/>
+      <circle cx="12" cy="12" r="9" />
+      <path d="M13 12h5" />
+      <path d="M13 15h4" />
+      <path d="M13 18h1" />
+      <path d="M13 9h4" />
+      <path d="M13 6h1" />
+    </svg>
+  </symbol>
+</svg>
+
+<input type="checkbox" class="sidebar-toggle" name="__navigation" id="__navigation">
+<input type="checkbox" class="sidebar-toggle" name="__toc" id="__toc">
+<label class="overlay sidebar-overlay" for="__navigation">
+  <div class="visually-hidden">Hide navigation sidebar</div>
+</label>
+<label class="overlay toc-overlay" for="__toc">
+  <div class="visually-hidden">Hide table of contents sidebar</div>
+</label>
+
+
+
+<div class="page">
+  <header class="mobile-header">
+    <div class="header-left">
+      <label class="nav-overlay-icon" for="__navigation">
+        <div class="visually-hidden">Toggle site navigation sidebar</div>
+        <i class="icon"><svg><use href="#svg-menu"></use></svg></i>
+      </label>
+    </div>
+    <div class="header-center">
+      <a href="index.html"><div class="brand">cyto-dl 0.4.0 documentation</div></a>
+    </div>
+    <div class="header-right">
+      <div class="theme-toggle-container theme-toggle-header">
+        <button class="theme-toggle">
+          <div class="visually-hidden">Toggle Light / Dark / Auto color theme</div>
+          <svg class="theme-icon-when-auto"><use href="#svg-sun-half"></use></svg>
+          <svg class="theme-icon-when-dark"><use href="#svg-moon"></use></svg>
+          <svg class="theme-icon-when-light"><use href="#svg-sun"></use></svg>
+        </button>
+      </div>
+      <label class="toc-overlay-icon toc-header-icon" for="__toc">
+        <div class="visually-hidden">Toggle table of contents sidebar</div>
+        <i class="icon"><svg><use href="#svg-toc"></use></svg></i>
+      </label>
+    </div>
+  </header>
+  <aside class="sidebar-drawer">
+    <div class="sidebar-container">
+      
+      <div class="sidebar-sticky"><a class="sidebar-brand" href="index.html">
+  
+  
+  <span class="sidebar-brand-text">cyto-dl 0.4.0 documentation</span>
+  
+</a><form class="sidebar-search-container" method="get" action="search.html" role="search">
+  <input class="sidebar-search" placeholder="Search" name="q" aria-label="Search">
+  <input type="hidden" name="check_keywords" value="yes">
+  <input type="hidden" name="area" value="default">
+</form>
+<div id="searchbox"></div><div class="sidebar-scroll"><div class="sidebar-tree">
+  <p class="caption" role="heading"><span class="caption-text">Contents:</span></p>
+<ul class="current">
+<li class="toctree-l1"><a class="reference internal" href="index.html">Home</a></li>
+<li class="toctree-l1"><a class="reference internal" href="installation.html">Installation</a></li>
+<li class="toctree-l1"><a class="reference internal" href="using_examples.html">Modifying the Example Configs</a></li>
+<li class="toctree-l1"><a class="reference internal" href="api.html">Using the API</a></li>
+<li class="toctree-l1 current has-children"><a class="reference internal" href="modules.html">Package modules</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-1" name="toctree-checkbox-1" role="switch" type="checkbox"/><label for="toctree-checkbox-1"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
+<li class="toctree-l2 current has-children"><a class="reference internal" href="cyto_dl.html">cyto_dl package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-2" name="toctree-checkbox-2" role="switch" type="checkbox"/><label for="toctree-checkbox-2"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.api.html">cyto_dl.api package</a><input class="toctree-checkbox" id="toctree-checkbox-3" name="toctree-checkbox-3" role="switch" type="checkbox"/><label for="toctree-checkbox-3"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.api.cyto_dl_model.html">cyto_dl.api.cyto_dl_model package</a><input class="toctree-checkbox" id="toctree-checkbox-4" name="toctree-checkbox-4" role="switch" type="checkbox"/><label for="toctree-checkbox-4"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.api.cyto_dl_model.cyto_dl_base_model.html">cyto_dl.api.cyto_dl_model.cyto_dl_base_model module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.api.cyto_dl_model.segmentation_plugin_model.html">cyto_dl.api.cyto_dl_model.segmentation_plugin_model module</a></li>
+</ul>
+</li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.api.data.html">cyto_dl.api.data module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.api.model.html">cyto_dl.api.model module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.callbacks.html">cyto_dl.callbacks package</a><input class="toctree-checkbox" id="toctree-checkbox-5" name="toctree-checkbox-5" role="switch" type="checkbox"/><label for="toctree-checkbox-5"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.callbacks.callback_utils.html">cyto_dl.callbacks.callback_utils module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.callbacks.csv_saver.html">cyto_dl.callbacks.csv_saver module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.callbacks.image_saver.html">cyto_dl.callbacks.image_saver module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.callbacks.latent_walk.html">cyto_dl.callbacks.latent_walk module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.callbacks.layer_freeze.html">cyto_dl.callbacks.layer_freeze module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.callbacks.model_utils.html">cyto_dl.callbacks.model_utils module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.callbacks.outlier_detection.html">cyto_dl.callbacks.outlier_detection module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.dataframe.html">cyto_dl.dataframe package</a><input class="toctree-checkbox" id="toctree-checkbox-6" name="toctree-checkbox-6" role="switch" type="checkbox"/><label for="toctree-checkbox-6"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.dataframe.transforms.html">cyto_dl.dataframe.transforms package</a><input class="toctree-checkbox" id="toctree-checkbox-7" name="toctree-checkbox-7" role="switch" type="checkbox"/><label for="toctree-checkbox-7"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.dataframe.transforms.filter.html">cyto_dl.dataframe.transforms.filter module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.dataframe.transforms.group_cols.html">cyto_dl.dataframe.transforms.group_cols module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.dataframe.transforms.misc.html">cyto_dl.dataframe.transforms.misc module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.dataframe.transforms.split.html">cyto_dl.dataframe.transforms.split module</a></li>
+</ul>
+</li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.dataframe.readers.html">cyto_dl.dataframe.readers module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.datamodules.html">cyto_dl.datamodules package</a><input class="toctree-checkbox" id="toctree-checkbox-8" name="toctree-checkbox-8" role="switch" type="checkbox"/><label for="toctree-checkbox-8"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.datamodules.dataframe.html">cyto_dl.datamodules.dataframe package</a><input class="toctree-checkbox" id="toctree-checkbox-9" name="toctree-checkbox-9" role="switch" type="checkbox"/><label for="toctree-checkbox-9"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.datamodules.dataframe.dataframe_datamodule.html">cyto_dl.datamodules.dataframe.dataframe_datamodule module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.datamodules.dataframe.grouped_dataframe_datamodule.html">cyto_dl.datamodules.dataframe.grouped_dataframe_datamodule module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.datamodules.dataframe.utils.html">cyto_dl.datamodules.dataframe.utils module</a></li>
+</ul>
+</li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.datamodules.array.html">cyto_dl.datamodules.array module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.datamodules.data_dict.html">cyto_dl.datamodules.data_dict module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.datamodules.folder.html">cyto_dl.datamodules.folder module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.datamodules.multidim_image.html">cyto_dl.datamodules.multidim_image module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.datamodules.smartcache.html">cyto_dl.datamodules.smartcache module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.datamodules.torchvision.html">cyto_dl.datamodules.torchvision module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.image.html">cyto_dl.image package</a><input class="toctree-checkbox" id="toctree-checkbox-10" name="toctree-checkbox-10" role="switch" type="checkbox"/><label for="toctree-checkbox-10"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.image.io.html">cyto_dl.image.io package</a><input class="toctree-checkbox" id="toctree-checkbox-11" name="toctree-checkbox-11" role="switch" type="checkbox"/><label for="toctree-checkbox-11"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.io.aicsimage_loader.html">cyto_dl.image.io.aicsimage_loader module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.io.monai_bio_reader.html">cyto_dl.image.io.monai_bio_reader module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.io.numpy_reader.html">cyto_dl.image.io.numpy_reader module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.io.ome_zarr_reader.html">cyto_dl.image.io.ome_zarr_reader module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.io.polygon_loader.html">cyto_dl.image.io.polygon_loader module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.io.skimage_reader.html">cyto_dl.image.io.skimage_reader module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.image.transforms.html">cyto_dl.image.transforms package</a><input class="toctree-checkbox" id="toctree-checkbox-12" name="toctree-checkbox-12" role="switch" type="checkbox"/><label for="toctree-checkbox-12"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.add_meta.html">cyto_dl.image.transforms.add_meta module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.bright_sampler.html">cyto_dl.image.transforms.bright_sampler module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.clip.html">cyto_dl.image.transforms.clip module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.contrastadjust.html">cyto_dl.image.transforms.contrastadjust module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.generate_jepa_masks.html">cyto_dl.image.transforms.generate_jepa_masks module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.merge.html">cyto_dl.image.transforms.merge module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.multiscale_cropper.html">cyto_dl.image.transforms.multiscale_cropper module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.pad.html">cyto_dl.image.transforms.pad module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.project.html">cyto_dl.image.transforms.project module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.rotation_mask_transform.html">cyto_dl.image.transforms.rotation_mask_transform module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.save.html">cyto_dl.image.transforms.save module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.so2_random_rotation.html">cyto_dl.image.transforms.so2_random_rotation module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.track_transforms.html">cyto_dl.image.transforms.track_transforms module</a></li>
+</ul>
+</li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.loggers.html">cyto_dl.loggers package</a><input class="toctree-checkbox" id="toctree-checkbox-13" name="toctree-checkbox-13" role="switch" type="checkbox"/><label for="toctree-checkbox-13"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.loggers.mlflow.html">cyto_dl.loggers.mlflow module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.models.html">cyto_dl.models package</a><input class="toctree-checkbox" id="toctree-checkbox-14" name="toctree-checkbox-14" role="switch" type="checkbox"/><label for="toctree-checkbox-14"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.models.classification.html">cyto_dl.models.classification package</a><input class="toctree-checkbox" id="toctree-checkbox-15" name="toctree-checkbox-15" role="switch" type="checkbox"/><label for="toctree-checkbox-15"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.classification.classification.html">cyto_dl.models.classification.classification module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.classification.timepoint_classification.html">cyto_dl.models.classification.timepoint_classification module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.models.contrastive.html">cyto_dl.models.contrastive package</a><input class="toctree-checkbox" id="toctree-checkbox-16" name="toctree-checkbox-16" role="switch" type="checkbox"/><label for="toctree-checkbox-16"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.contrastive.contrastive.html">cyto_dl.models.contrastive.contrastive module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.models.handlers.html">cyto_dl.models.handlers package</a><input class="toctree-checkbox" id="toctree-checkbox-17" name="toctree-checkbox-17" role="switch" type="checkbox"/><label for="toctree-checkbox-17"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.handlers.base_handler.html">cyto_dl.models.handlers.base_handler module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.handlers.image_handler.html">cyto_dl.models.handlers.image_handler module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.handlers.load_image_patch.html">cyto_dl.models.handlers.load_image_patch module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.models.im2im.html">cyto_dl.models.im2im package</a><input class="toctree-checkbox" id="toctree-checkbox-18" name="toctree-checkbox-18" role="switch" type="checkbox"/><label for="toctree-checkbox-18"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.models.im2im.utils.html">cyto_dl.models.im2im.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-19" name="toctree-checkbox-19" role="switch" type="checkbox"/><label for="toctree-checkbox-19"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.models.im2im.utils.postprocessing.html">cyto_dl.models.im2im.utils.postprocessing package</a><input class="toctree-checkbox" id="toctree-checkbox-20" name="toctree-checkbox-20" role="switch" type="checkbox"/><label for="toctree-checkbox-20"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.models.im2im.utils.postprocessing.act_thresh_label.html">cyto_dl.models.im2im.utils.postprocessing.act_thresh_label module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.models.im2im.utils.postprocessing.auto_thresh.html">cyto_dl.models.im2im.utils.postprocessing.auto_thresh module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.models.im2im.utils.postprocessing.dict_to_im.html">cyto_dl.models.im2im.utils.postprocessing.dict_to_im module</a></li>
+</ul>
+</li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.models.im2im.utils.instance_seg.html">cyto_dl.models.im2im.utils.instance_seg module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.models.im2im.utils.noise_annealer.html">cyto_dl.models.im2im.utils.noise_annealer module</a></li>
+</ul>
+</li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.im2im.gan.html">cyto_dl.models.im2im.gan module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.im2im.multi_task.html">cyto_dl.models.im2im.multi_task module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.models.jepa.html">cyto_dl.models.jepa package</a><input class="toctree-checkbox" id="toctree-checkbox-21" name="toctree-checkbox-21" role="switch" type="checkbox"/><label for="toctree-checkbox-21"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.jepa.ijepa.html">cyto_dl.models.jepa.ijepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.jepa.iwm.html">cyto_dl.models.jepa.iwm module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.jepa.jepa_base.html">cyto_dl.models.jepa.jepa_base module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.models.utils.html">cyto_dl.models.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-22" name="toctree-checkbox-22" role="switch" type="checkbox"/><label for="toctree-checkbox-22"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.utils.mlflow.html">cyto_dl.models.utils.mlflow module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.utils.utils.html">cyto_dl.models.utils.utils module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.models.vae.html">cyto_dl.models.vae package</a><input class="toctree-checkbox" id="toctree-checkbox-23" name="toctree-checkbox-23" role="switch" type="checkbox"/><label for="toctree-checkbox-23"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.models.vae.o2_spharm_vae.html">cyto_dl.models.vae.o2_spharm_vae package</a><input class="toctree-checkbox" id="toctree-checkbox-24" name="toctree-checkbox-24" role="switch" type="checkbox"/><label for="toctree-checkbox-24"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.models.vae.o2_spharm_vae.o2_spharm_encoder.html">cyto_dl.models.vae.o2_spharm_vae.o2_spharm_encoder module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.models.vae.o2_spharm_vae.o2_spharm_vae.html">cyto_dl.models.vae.o2_spharm_vae.o2_spharm_vae module</a></li>
+</ul>
+</li>
+<li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.models.vae.priors.html">cyto_dl.models.vae.priors package</a><input class="toctree-checkbox" id="toctree-checkbox-25" name="toctree-checkbox-25" role="switch" type="checkbox"/><label for="toctree-checkbox-25"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.models.vae.priors.abstract_prior.html">cyto_dl.models.vae.priors.abstract_prior module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.models.vae.priors.gaussian.html">cyto_dl.models.vae.priors.gaussian module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.models.vae.priors.identity_prior.html">cyto_dl.models.vae.priors.identity_prior module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.models.vae.priors.joint_prior.html">cyto_dl.models.vae.priors.joint_prior module</a></li>
+</ul>
+</li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.vae.base_vae.html">cyto_dl.models.vae.base_vae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.vae.image_canon_vae.html">cyto_dl.models.vae.image_canon_vae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.vae.image_encoder.html">cyto_dl.models.vae.image_encoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.vae.image_vae.html">cyto_dl.models.vae.image_vae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.vae.latent_loss_vae.html">cyto_dl.models.vae.latent_loss_vae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.vae.point_cloud_vae.html">cyto_dl.models.vae.point_cloud_vae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.vae.tabular_vae.html">cyto_dl.models.vae.tabular_vae module</a></li>
+</ul>
+</li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.models.base_model.html">cyto_dl.models.base_model module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.models.basic_model.html">cyto_dl.models.basic_model module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 current has-children"><a class="reference internal" href="cyto_dl.nn.html">cyto_dl.nn package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-26" name="toctree-checkbox-26" role="switch" type="checkbox"/><label for="toctree-checkbox-26"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.discriminators.html">cyto_dl.nn.discriminators package</a><input class="toctree-checkbox" id="toctree-checkbox-27" name="toctree-checkbox-27" role="switch" type="checkbox"/><label for="toctree-checkbox-27"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.discriminators.multi_scale_discriminator.html">cyto_dl.nn.discriminators.multi_scale_discriminator module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.discriminators.n_layer_discriminator.html">cyto_dl.nn.discriminators.n_layer_discriminator module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.head.html">cyto_dl.nn.head package</a><input class="toctree-checkbox" id="toctree-checkbox-28" name="toctree-checkbox-28" role="switch" type="checkbox"/><label for="toctree-checkbox-28"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.head.base_head.html">cyto_dl.nn.head.base_head module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.head.gan_head.html">cyto_dl.nn.head.gan_head module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.head.gan_head_superres.html">cyto_dl.nn.head.gan_head_superres module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.head.mae_head.html">cyto_dl.nn.head.mae_head module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.head.mask_head.html">cyto_dl.nn.head.mask_head module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.head.res_blocks_head.html">cyto_dl.nn.head.res_blocks_head module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.head.vic_reg.html">cyto_dl.nn.head.vic_reg module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.losses.html">cyto_dl.nn.losses package</a><input class="toctree-checkbox" id="toctree-checkbox-29" name="toctree-checkbox-29" role="switch" type="checkbox"/><label for="toctree-checkbox-29"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.adversarial_loss.html">cyto_dl.nn.losses.adversarial_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.chamfer_loss.html">cyto_dl.nn.losses.chamfer_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.continuous_bernoulli.html">cyto_dl.nn.losses.continuous_bernoulli module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.cosine_loss.html">cyto_dl.nn.losses.cosine_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.gan_loss.html">cyto_dl.nn.losses.gan_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.gaussian_nll_loss.html">cyto_dl.nn.losses.gaussian_nll_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.geomloss.html">cyto_dl.nn.losses.geomloss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.loss_wrapper.html">cyto_dl.nn.losses.loss_wrapper module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.spharm_loss.html">cyto_dl.nn.losses.spharm_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.threshold_loss.html">cyto_dl.nn.losses.threshold_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.vic_reg.html">cyto_dl.nn.losses.vic_reg module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.weibull.html">cyto_dl.nn.losses.weibull module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.weighted_mse_loss.html">cyto_dl.nn.losses.weighted_mse_loss module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.point_cloud.html">cyto_dl.nn.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-30" name="toctree-checkbox-30" role="switch" type="checkbox"/><label for="toctree-checkbox-30"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.point_cloud.dgcnn.html">cyto_dl.nn.point_cloud.dgcnn module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.point_cloud.folding_net.html">cyto_dl.nn.point_cloud.folding_net module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.point_cloud.graph_functions.html">cyto_dl.nn.point_cloud.graph_functions module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.point_cloud.vnn.html">cyto_dl.nn.point_cloud.vnn module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 current has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
+<li class="toctree-l5 current has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
+<li class="toctree-l6 current current-page"><a class="current reference internal" href="#">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
+</ul>
+</li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
+</ul>
+</li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.hr_skip.html">cyto_dl.nn.hr_skip module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.mlp.html">cyto_dl.nn.mlp module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.res_unit.html">cyto_dl.nn.res_unit module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.spatial_transformer.html">cyto_dl.nn.spatial_transformer module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
+</ul>
+</li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
+</ul>
+</li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.utils.array.html">cyto_dl.utils.array module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.utils.checkpoint.html">cyto_dl.utils.checkpoint module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.utils.config.html">cyto_dl.utils.config module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.utils.download_test_data.html">cyto_dl.utils.download_test_data module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.utils.dummy_dataset.html">cyto_dl.utils.dummy_dataset module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.utils.pylogger.html">cyto_dl.utils.pylogger module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.utils.rich_utils.html">cyto_dl.utils.rich_utils module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.utils.rotation.html">cyto_dl.utils.rotation module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.utils.template_utils.html">cyto_dl.utils.template_utils module</a></li>
+</ul>
+</li>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.compile.html">cyto_dl.compile module</a></li>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.eval.html">cyto_dl.eval module</a></li>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.train.html">cyto_dl.train module</a></li>
+</ul>
+</li>
+</ul>
+</li>
+</ul>
+
+</div>
+</div>
+
+      </div>
+      
+    </div>
+  </aside>
+  <div class="main">
+    <div class="content">
+      <div class="article-container">
+        <a href="#" class="back-to-top muted-link">
+          <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24">
+            <path d="M13 20h-2V8l-5.5 5.5-1.42-1.42L12 4.16l7.92 7.92-1.42 1.42L13 8v12z"></path>
+          </svg>
+          <span>Back to top</span>
+        </a>
+        <div class="content-icon-container">
+          
+<div class="theme-toggle-container theme-toggle-content">
+            <button class="theme-toggle">
+              <div class="visually-hidden">Toggle Light / Dark / Auto color theme</div>
+              <svg class="theme-icon-when-auto"><use href="#svg-sun-half"></use></svg>
+              <svg class="theme-icon-when-dark"><use href="#svg-moon"></use></svg>
+              <svg class="theme-icon-when-light"><use href="#svg-sun"></use></svg>
+            </button>
+          </div>
+          <label class="toc-overlay-icon toc-content-icon" for="__toc">
+            <div class="visually-hidden">Toggle table of contents sidebar</div>
+            <i class="icon"><svg><use href="#svg-toc"></use></svg></i>
+          </label>
+        </div>
+        <article role="main">
+          <section id="module-cyto_dl.nn.vits.blocks.masked_unit_attention">
+<span id="cyto-dl-nn-vits-blocks-masked-unit-attention-module"></span><h1>cyto_dl.nn.vits.blocks.masked_unit_attention module<a class="headerlink" href="#module-cyto_dl.nn.vits.blocks.masked_unit_attention" title="Permalink to this heading">#</a></h1>
+<dl class="py class">
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.blocks.masked_unit_attention.HieraBlock">
+<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">cyto_dl.nn.vits.blocks.masked_unit_attention.</span></span><span class="sig-name descname"><span class="pre">HieraBlock</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="pre">dim:</span> <span class="pre">int,</span> <span class="pre">dim_out:</span> <span class="pre">int,</span> <span class="pre">heads:</span> <span class="pre">int,</span> <span class="pre">spatial_dims:</span> <span class="pre">int</span> <span class="pre">=</span> <span class="pre">3,</span> <span class="pre">mlp_ratio:</span> <span class="pre">float</span> <span class="pre">=</span> <span class="pre">4.0,</span> <span class="pre">drop_path:</span> <span class="pre">float</span> <span class="pre">=</span> <span class="pre">0.0,</span> <span class="pre">norm_layer:</span> <span class="pre">~torch.nn.modules.module.Module</span> <span class="pre">=</span> <span class="pre">&lt;class</span> <span class="pre">'torch.nn.modules.normalization.LayerNorm'&gt;,</span> <span class="pre">act_layer:</span> <span class="pre">~torch.nn.modules.module.Module</span> <span class="pre">=</span> <span class="pre">&lt;class</span> <span class="pre">'torch.nn.modules.activation.GELU'&gt;,</span> <span class="pre">q_stride:</span> <span class="pre">~typing.List[int]</span> <span class="pre">=</span> <span class="pre">[1,</span> <span class="pre">1,</span> <span class="pre">1],</span> <span class="pre">patches_per_mask_unit:</span> <span class="pre">~typing.List[int]</span> <span class="pre">=</span> <span class="pre">[2,</span> <span class="pre">12,</span> <span class="pre">12]</span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/blocks/masked_unit_attention.html#HieraBlock"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.blocks.masked_unit_attention.HieraBlock" title="Permalink to this definition">#</a></dt>
+<dd><p>Bases: <code class="xref py py-class docutils literal notranslate"><span class="pre">Module</span></code></p>
+<dl class="field-list simple">
+<dt class="field-odd">Parameters<span class="colon">:</span></dt>
+<dd class="field-odd"><ul class="simple">
+<li><p><strong>dim</strong> (<em>int</em>) – Dimension of the input features.</p></li>
+<li><p><strong>dim_out</strong> (<em>int</em>) – Dimension of the output features.</p></li>
+<li><p><strong>heads</strong> (<em>int</em>) – Number of attention heads.</p></li>
+<li><p><strong>spatial_dims</strong> (<em>int, optional</em>) – Number of spatial dimensions, by default 3.</p></li>
+<li><p><strong>mlp_ratio</strong> (<em>float, optional</em>) – Ratio of MLP hidden dim to embedding dim, by default 4.0.</p></li>
+<li><p><strong>drop_path</strong> (<em>float, optional</em>) – Dropout rate for the path, by default 0.0.</p></li>
+<li><p><strong>norm_layer</strong> (<em>nn.Module, optional</em>) – Normalization layer, by default nn.LayerNorm.</p></li>
+<li><p><strong>act_layer</strong> (<em>nn.Module, optional</em>) – Activation layer for the MLP, by default nn.GELU.</p></li>
+<li><p><strong>q_stride</strong> (<em>List[int], optional</em>) – Stride for query, by default [1, 1, 1].</p></li>
+<li><p><strong>patches_per_mask_unit</strong> (<em>List[int], optional</em>) – Number of patches per mask unit, by default [2, 12, 12].</p></li>
+</ul>
+</dd>
+</dl>
+<dl class="py method">
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.blocks.masked_unit_attention.HieraBlock.forward">
+<span class="sig-name descname"><span class="pre">forward</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">x</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">Tensor</span></span></em><span class="sig-paren">)</span> <span class="sig-return"><span class="sig-return-icon">&#x2192;</span> <span class="sig-return-typehint"><span class="pre">Tensor</span></span></span><a class="reference internal" href="_modules/cyto_dl/nn/vits/blocks/masked_unit_attention.html#HieraBlock.forward"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.blocks.masked_unit_attention.HieraBlock.forward" title="Permalink to this definition">#</a></dt>
+<dd><p>x: batch x mask units x tokens x emb_dim</p>
+</dd></dl>
+
+</dd></dl>
+
+<dl class="py class">
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.blocks.masked_unit_attention.MaskUnitAttention">
+<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">cyto_dl.nn.vits.blocks.masked_unit_attention.</span></span><span class="sig-name descname"><span class="pre">MaskUnitAttention</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">dim</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">dim_out</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">spatial_dims</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">3</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">num_heads</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">8</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">qkv_bias</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">False</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">attn_drop</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">0.0</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">proj_drop</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">0.0</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">q_stride</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">[1,</span> <span class="pre">1,</span> <span class="pre">1]</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">patches_per_mask_unit</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">[2,</span> <span class="pre">12,</span> <span class="pre">12]</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/blocks/masked_unit_attention.html#MaskUnitAttention"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.blocks.masked_unit_attention.MaskUnitAttention" title="Permalink to this definition">#</a></dt>
+<dd><p>Bases: <code class="xref py py-class docutils literal notranslate"><span class="pre">Module</span></code></p>
+<dl class="field-list simple">
+<dt class="field-odd">Parameters<span class="colon">:</span></dt>
+<dd class="field-odd"><ul class="simple">
+<li><p><strong>dim</strong> (<em>int</em>) – Dimension of the input features.</p></li>
+<li><p><strong>dim_out</strong> (<em>int</em>) – Dimension of the output features.</p></li>
+<li><p><strong>spatial_dims</strong> (<em>int, optional</em>) – Number of spatial dimensions, by default 3.</p></li>
+<li><p><strong>num_heads</strong> (<em>int, optional</em>) – Number of attention heads, by default 8.</p></li>
+<li><p><strong>qkv_bias</strong> (<em>bool, optional</em>) – If True, add a learnable bias to query, key, value, by default False.</p></li>
+<li><p><strong>attn_drop</strong> (<em>float, optional</em>) – Dropout rate for attention, by default 0.0.</p></li>
+<li><p><strong>proj_drop</strong> (<em>float, optional</em>) – Dropout rate for projection, by default 0.0.</p></li>
+<li><p><strong>q_stride</strong> (<em>List[int], optional</em>) – Stride for query, by default [1, 1, 1].</p></li>
+<li><p><strong>patches_per_mask_unit</strong> (<em>List[int], optional</em>) – Number of patches per mask unit, by default [2, 12, 12].</p></li>
+</ul>
+</dd>
+</dl>
+<dl class="py method">
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.blocks.masked_unit_attention.MaskUnitAttention.forward">
+<span class="sig-name descname"><span class="pre">forward</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">x</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/blocks/masked_unit_attention.html#MaskUnitAttention.forward"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.blocks.masked_unit_attention.MaskUnitAttention.forward" title="Permalink to this definition">#</a></dt>
+<dd></dd></dl>
+
+</dd></dl>
+
+</section>
+
+        </article>
+      </div>
+      <footer>
+        
+        <div class="related-pages">
+          <a class="next-page" href="cyto_dl.nn.vits.decoder.html">
+              <div class="page-info">
+                <div class="context">
+                  <span>Next</span>
+                </div>
+                <div class="title">cyto_dl.nn.vits.decoder module</div>
+              </div>
+              <svg class="furo-related-icon"><use href="#svg-arrow-right"></use></svg>
+            </a>
+          <a class="prev-page" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">
+              <svg class="furo-related-icon"><use href="#svg-arrow-right"></use></svg>
+              <div class="page-info">
+                <div class="context">
+                  <span>Previous</span>
+                </div>
+                
+                <div class="title">cyto_dl.nn.vits.blocks.intermediate_weigher module</div>
+                
+              </div>
+            </a>
+        </div>
+        <div class="bottom-of-page">
+          <div class="left-details">
+            <div class="copyright">
+                Copyright &#169; 2023
+            </div>
+            Made with <a href="https://www.sphinx-doc.org/">Sphinx</a> and <a class="muted-link" href="https://pradyunsg.me">@pradyunsg</a>'s
+            
+            <a href="https://github.com/pradyunsg/furo">Furo</a>
+            
+          </div>
+          <div class="right-details">
+            
+          </div>
+        </div>
+        
+      </footer>
+    </div>
+    <aside class="toc-drawer">
+      
+      
+      <div class="toc-sticky toc-scroll">
+        <div class="toc-title-container">
+          <span class="toc-title">
+            On this page
+          </span>
+        </div>
+        <div class="toc-tree-container">
+          <div class="toc-tree">
+            <ul>
+<li><a class="reference internal" href="#">cyto_dl.nn.vits.blocks.masked_unit_attention module</a><ul>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.blocks.masked_unit_attention.HieraBlock"><code class="docutils literal notranslate"><span class="pre">HieraBlock</span></code></a><ul>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.blocks.masked_unit_attention.HieraBlock.forward"><code class="docutils literal notranslate"><span class="pre">HieraBlock.forward()</span></code></a></li>
+</ul>
+</li>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.blocks.masked_unit_attention.MaskUnitAttention"><code class="docutils literal notranslate"><span class="pre">MaskUnitAttention</span></code></a><ul>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.blocks.masked_unit_attention.MaskUnitAttention.forward"><code class="docutils literal notranslate"><span class="pre">MaskUnitAttention.forward()</span></code></a></li>
+</ul>
+</li>
+</ul>
+</li>
+</ul>
+
+          </div>
+        </div>
+      </div>
+      
+      
+    </aside>
+  </div>
+</div><script data-url_root="./" id="documentation_options" src="_static/documentation_options.js"></script>
+    <script src="_static/doctools.js"></script>
+    <script src="_static/sphinx_highlight.js"></script>
+    <script src="_static/scripts/furo.js"></script>
+    </body>
+</html>
\ No newline at end of file
diff --git a/cyto_dl.nn.vits.blocks.patchify.html b/cyto_dl.nn.vits.blocks.patchify.html
index c85bbda8c..fdc2db069 100644
--- a/cyto_dl.nn.vits.blocks.patchify.html
+++ b/cyto_dl.nn.vits.blocks.patchify.html
@@ -3,10 +3,10 @@
   <head><meta charset="utf-8"/>
     <meta name="viewport" content="width=device-width,initial-scale=1"/>
     <meta name="color-scheme" content="light dark"><meta name="generator" content="Docutils 0.19: https://docutils.sourceforge.io/" />
-<link rel="index" title="Index" href="genindex.html" /><link rel="search" title="Search" href="search.html" /><link rel="next" title="cyto_dl.nn.vits.cross_mae module" href="cyto_dl.nn.vits.cross_mae.html" /><link rel="prev" title="cyto_dl.nn.vits.blocks.intermediate_weigher module" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html" />
+<link rel="index" title="Index" href="genindex.html" /><link rel="search" title="Search" href="search.html" /><link rel="next" title="cyto_dl.nn.vits.blocks.patchify.patchify module" href="cyto_dl.nn.vits.blocks.patchify.patchify.html" /><link rel="prev" title="cyto_dl.nn.vits.blocks package" href="cyto_dl.nn.vits.blocks.html" />
 
     <!-- Generated with Sphinx 6.1.3 and Furo 2023.03.23 -->
-        <title>cyto_dl.nn.vits.blocks.patchify module - cyto-dl 0.4.0 documentation</title>
+        <title>cyto_dl.nn.vits.blocks.patchify package - cyto-dl 0.4.0 documentation</title>
       <link rel="stylesheet" type="text/css" href="_static/pygments.css" />
     <link rel="stylesheet" type="text/css" href="_static/styles/furo.css?digest=fad236701ea90a88636c2a8c73b44ae642ed2a53" />
     <link rel="stylesheet" type="text/css" href="_static/styles/furo-extensions.css?digest=30d1aed668e5c3a91c3e3bf6a60b675221979f0e" />
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 current has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
 <li class="toctree-l5 current has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
+<li class="toctree-l6 current has-children current-page"><a class="current reference internal" href="#">cyto_dl.nn.vits.blocks.patchify package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6 current current-page"><a class="current reference internal" href="#">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
@@ -435,43 +442,48 @@
         </div>
         <article role="main">
           <section id="module-cyto_dl.nn.vits.blocks.patchify">
-<span id="cyto-dl-nn-vits-blocks-patchify-module"></span><h1>cyto_dl.nn.vits.blocks.patchify module<a class="headerlink" href="#module-cyto_dl.nn.vits.blocks.patchify" title="Permalink to this heading">#</a></h1>
-<dl class="py class">
-<dt class="sig sig-object py" id="cyto_dl.nn.vits.blocks.patchify.Patchify">
-<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">cyto_dl.nn.vits.blocks.patchify.</span></span><span class="sig-name descname"><span class="pre">Patchify</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">patch_size</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">]</span></span></span></em>, <em class="sig-param"><span class="n"><span class="pre">emb_dim</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">n_patches</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">]</span></span></span></em>, <em class="sig-param"><span class="n"><span class="pre">spatial_dims</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">3</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">context_pixels</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">]</span></span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">[0,</span> <span class="pre">0,</span> <span class="pre">0]</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">input_channels</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">1</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">tasks</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">str</span><span class="p"><span class="pre">]</span></span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">[]</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">learnable_pos_embedding</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">bool</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">True</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/blocks/patchify.html#Patchify"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.blocks.patchify.Patchify" title="Permalink to this definition">#</a></dt>
-<dd><p>Bases: <code class="xref py py-class docutils literal notranslate"><span class="pre">Module</span></code></p>
-<p>Class for converting images to a masked sequence of patches with positional embeddings.</p>
-<dl class="field-list simple">
-<dt class="field-odd">Parameters<span class="colon">:</span></dt>
-<dd class="field-odd"><ul class="simple">
-<li><p><strong>patch_size</strong> (<em>List[int]</em>) – Size of each patch in pix (ZYX order for 3D, YX order for 2D)</p></li>
-<li><p><strong>emb_dim</strong> (<em>int</em>) – Dimension of encoder</p></li>
-<li><p><strong>n_patches</strong> (<em>List[int]</em>) – Number of patches in each spatial dimension (ZYX order for 3D, YX order for 2D)</p></li>
-<li><p><strong>spatial_dims</strong> (<em>int</em>) – Number of spatial dimensions</p></li>
-<li><p><strong>context_pixels</strong> (<em>List[int]</em>) – Number of extra pixels around each patch to include in convolutional embedding to encoder dimension.</p></li>
-<li><p><strong>input_channels</strong> (<em>int</em>) – Number of input channels</p></li>
-<li><p><strong>tasks</strong> (<em>List[str]</em>) – List of tasks to encode</p></li>
-<li><p><strong>learnable_pos_embedding</strong> (<em>bool</em>) – If True, learnable positional embeddings are used. If False, fixed sin/cos positional embeddings. Empirically, fixed positional embeddings work better for brightfield images.</p></li>
-</ul>
-</dd>
-</dl>
-<dl class="py method">
-<dt class="sig sig-object py" id="cyto_dl.nn.vits.blocks.patchify.Patchify.forward">
-<span class="sig-name descname"><span class="pre">forward</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">img</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">mask_ratio</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">task</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/blocks/patchify.html#Patchify.forward"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.blocks.patchify.Patchify.forward" title="Permalink to this definition">#</a></dt>
-<dd></dd></dl>
-
-<dl class="py method">
-<dt class="sig sig-object py" id="cyto_dl.nn.vits.blocks.patchify.Patchify.get_mask">
-<span class="sig-name descname"><span class="pre">get_mask</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">img</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">n_visible_patches</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">num_patches</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/blocks/patchify.html#Patchify.get_mask"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.blocks.patchify.Patchify.get_mask" title="Permalink to this definition">#</a></dt>
-<dd></dd></dl>
-
-</dd></dl>
-
-<dl class="py function">
-<dt class="sig sig-object py" id="cyto_dl.nn.vits.blocks.patchify.random_indexes">
-<span class="sig-prename descclassname"><span class="pre">cyto_dl.nn.vits.blocks.patchify.</span></span><span class="sig-name descname"><span class="pre">random_indexes</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">size</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">device</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/blocks/patchify.html#random_indexes"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.blocks.patchify.random_indexes" title="Permalink to this definition">#</a></dt>
-<dd></dd></dl>
-
+<span id="cyto-dl-nn-vits-blocks-patchify-package"></span><h1>cyto_dl.nn.vits.blocks.patchify package<a class="headerlink" href="#module-cyto_dl.nn.vits.blocks.patchify" title="Permalink to this heading">#</a></h1>
+<section id="submodules">
+<h2>Submodules<a class="headerlink" href="#submodules" title="Permalink to this heading">#</a></h2>
+<div class="toctree-wrapper compound">
+<ul>
+<li class="toctree-l1"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a><ul>
+<li class="toctree-l2"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html#cyto_dl.nn.vits.blocks.patchify.patchify.Patchify"><code class="docutils literal notranslate"><span class="pre">Patchify</span></code></a><ul>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html#cyto_dl.nn.vits.blocks.patchify.patchify.Patchify.create_img2token"><code class="docutils literal notranslate"><span class="pre">Patchify.create_img2token()</span></code></a></li>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html#cyto_dl.nn.vits.blocks.patchify.patchify.Patchify.extract_visible_tokens"><code class="docutils literal notranslate"><span class="pre">Patchify.extract_visible_tokens()</span></code></a></li>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html#cyto_dl.nn.vits.blocks.patchify.patchify.Patchify.get_mask_args"><code class="docutils literal notranslate"><span class="pre">Patchify.get_mask_args()</span></code></a></li>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html#cyto_dl.nn.vits.blocks.patchify.patchify.Patchify.img2token"><code class="docutils literal notranslate"><span class="pre">Patchify.img2token</span></code></a></li>
+</ul>
+</li>
+</ul>
+</li>
+<li class="toctree-l1"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a><ul>
+<li class="toctree-l2"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html#cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase"><code class="docutils literal notranslate"><span class="pre">PatchifyBase</span></code></a><ul>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html#cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase.create_conv"><code class="docutils literal notranslate"><span class="pre">PatchifyBase.create_conv()</span></code></a></li>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html#cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase.create_patch2img"><code class="docutils literal notranslate"><span class="pre">PatchifyBase.create_patch2img()</span></code></a></li>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html#cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase.extract_visible_tokens"><code class="docutils literal notranslate"><span class="pre">PatchifyBase.extract_visible_tokens()</span></code></a></li>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html#cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase.forward"><code class="docutils literal notranslate"><span class="pre">PatchifyBase.forward()</span></code></a></li>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html#cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase.get_mask"><code class="docutils literal notranslate"><span class="pre">PatchifyBase.get_mask()</span></code></a></li>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html#cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase.get_mask_args"><code class="docutils literal notranslate"><span class="pre">PatchifyBase.get_mask_args()</span></code></a></li>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html#cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase.img2token"><code class="docutils literal notranslate"><span class="pre">PatchifyBase.img2token</span></code></a></li>
+</ul>
+</li>
+</ul>
+</li>
+<li class="toctree-l1"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a><ul>
+<li class="toctree-l2"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html#cyto_dl.nn.vits.blocks.patchify.patchify_hiera.PatchifyHiera"><code class="docutils literal notranslate"><span class="pre">PatchifyHiera</span></code></a><ul>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html#cyto_dl.nn.vits.blocks.patchify.patchify_hiera.PatchifyHiera.create_img2token"><code class="docutils literal notranslate"><span class="pre">PatchifyHiera.create_img2token()</span></code></a></li>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html#cyto_dl.nn.vits.blocks.patchify.patchify_hiera.PatchifyHiera.extract_visible_tokens"><code class="docutils literal notranslate"><span class="pre">PatchifyHiera.extract_visible_tokens()</span></code></a></li>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html#cyto_dl.nn.vits.blocks.patchify.patchify_hiera.PatchifyHiera.get_mask_args"><code class="docutils literal notranslate"><span class="pre">PatchifyHiera.get_mask_args()</span></code></a></li>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html#cyto_dl.nn.vits.blocks.patchify.patchify_hiera.PatchifyHiera.img2token"><code class="docutils literal notranslate"><span class="pre">PatchifyHiera.img2token</span></code></a></li>
+</ul>
+</li>
+<li class="toctree-l2"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html#cyto_dl.nn.vits.blocks.patchify.patchify_hiera.take_indexes_mask"><code class="docutils literal notranslate"><span class="pre">take_indexes_mask()</span></code></a></li>
+</ul>
+</li>
+</ul>
+</div>
+</section>
 </section>
 
         </article>
@@ -479,23 +491,23 @@
       <footer>
         
         <div class="related-pages">
-          <a class="next-page" href="cyto_dl.nn.vits.cross_mae.html">
+          <a class="next-page" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">
               <div class="page-info">
                 <div class="context">
                   <span>Next</span>
                 </div>
-                <div class="title">cyto_dl.nn.vits.cross_mae module</div>
+                <div class="title">cyto_dl.nn.vits.blocks.patchify.patchify module</div>
               </div>
               <svg class="furo-related-icon"><use href="#svg-arrow-right"></use></svg>
             </a>
-          <a class="prev-page" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">
+          <a class="prev-page" href="cyto_dl.nn.vits.blocks.html">
               <svg class="furo-related-icon"><use href="#svg-arrow-right"></use></svg>
               <div class="page-info">
                 <div class="context">
                   <span>Previous</span>
                 </div>
                 
-                <div class="title">cyto_dl.nn.vits.blocks.intermediate_weigher module</div>
+                <div class="title">cyto_dl.nn.vits.blocks package</div>
                 
               </div>
             </a>
@@ -529,13 +541,8 @@
         <div class="toc-tree-container">
           <div class="toc-tree">
             <ul>
-<li><a class="reference internal" href="#">cyto_dl.nn.vits.blocks.patchify module</a><ul>
-<li><a class="reference internal" href="#cyto_dl.nn.vits.blocks.patchify.Patchify"><code class="docutils literal notranslate"><span class="pre">Patchify</span></code></a><ul>
-<li><a class="reference internal" href="#cyto_dl.nn.vits.blocks.patchify.Patchify.forward"><code class="docutils literal notranslate"><span class="pre">Patchify.forward()</span></code></a></li>
-<li><a class="reference internal" href="#cyto_dl.nn.vits.blocks.patchify.Patchify.get_mask"><code class="docutils literal notranslate"><span class="pre">Patchify.get_mask()</span></code></a></li>
-</ul>
-</li>
-<li><a class="reference internal" href="#cyto_dl.nn.vits.blocks.patchify.random_indexes"><code class="docutils literal notranslate"><span class="pre">random_indexes()</span></code></a></li>
+<li><a class="reference internal" href="#">cyto_dl.nn.vits.blocks.patchify package</a><ul>
+<li><a class="reference internal" href="#submodules">Submodules</a></li>
 </ul>
 </li>
 </ul>
diff --git a/cyto_dl.nn.vits.cross_mae.html b/cyto_dl.nn.vits.blocks.patchify.patchify.html
similarity index 80%
rename from cyto_dl.nn.vits.cross_mae.html
rename to cyto_dl.nn.vits.blocks.patchify.patchify.html
index 866aeba94..a732034ab 100644
--- a/cyto_dl.nn.vits.cross_mae.html
+++ b/cyto_dl.nn.vits.blocks.patchify.patchify.html
@@ -3,10 +3,10 @@
   <head><meta charset="utf-8"/>
     <meta name="viewport" content="width=device-width,initial-scale=1"/>
     <meta name="color-scheme" content="light dark"><meta name="generator" content="Docutils 0.19: https://docutils.sourceforge.io/" />
-<link rel="index" title="Index" href="genindex.html" /><link rel="search" title="Search" href="search.html" /><link rel="next" title="cyto_dl.nn.vits.jepa module" href="cyto_dl.nn.vits.jepa.html" /><link rel="prev" title="cyto_dl.nn.vits.blocks.patchify module" href="cyto_dl.nn.vits.blocks.patchify.html" />
+<link rel="index" title="Index" href="genindex.html" /><link rel="search" title="Search" href="search.html" /><link rel="next" title="cyto_dl.nn.vits.blocks.patchify.patchify_base module" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html" /><link rel="prev" title="cyto_dl.nn.vits.blocks.patchify package" href="cyto_dl.nn.vits.blocks.patchify.html" />
 
     <!-- Generated with Sphinx 6.1.3 and Furo 2023.03.23 -->
-        <title>cyto_dl.nn.vits.cross_mae module - cyto-dl 0.4.0 documentation</title>
+        <title>cyto_dl.nn.vits.blocks.patchify.patchify module - cyto-dl 0.4.0 documentation</title>
       <link rel="stylesheet" type="text/css" href="_static/pygments.css" />
     <link rel="stylesheet" type="text/css" href="_static/styles/furo.css?digest=fad236701ea90a88636c2a8c73b44ae642ed2a53" />
     <link rel="stylesheet" type="text/css" href="_static/styles/furo-extensions.css?digest=30d1aed668e5c3a91c3e3bf6a60b675221979f0e" />
@@ -349,15 +349,22 @@
 </ul>
 </li>
 <li class="toctree-l4 current has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
-<li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5 current has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
+<li class="toctree-l6 current has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
+<li class="toctree-l7 current current-page"><a class="current reference internal" href="#">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5 current current-page"><a class="current reference internal" href="#">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
@@ -434,35 +441,32 @@
           </label>
         </div>
         <article role="main">
-          <section id="module-cyto_dl.nn.vits.cross_mae">
-<span id="cyto-dl-nn-vits-cross-mae-module"></span><h1>cyto_dl.nn.vits.cross_mae module<a class="headerlink" href="#module-cyto_dl.nn.vits.cross_mae" title="Permalink to this heading">#</a></h1>
+          <section id="module-cyto_dl.nn.vits.blocks.patchify.patchify">
+<span id="cyto-dl-nn-vits-blocks-patchify-patchify-module"></span><h1>cyto_dl.nn.vits.blocks.patchify.patchify module<a class="headerlink" href="#module-cyto_dl.nn.vits.blocks.patchify.patchify" title="Permalink to this heading">#</a></h1>
 <dl class="py class">
-<dt class="sig sig-object py" id="cyto_dl.nn.vits.cross_mae.CrossMAE_Decoder">
-<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">cyto_dl.nn.vits.cross_mae.</span></span><span class="sig-name descname"><span class="pre">CrossMAE_Decoder</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">num_patches</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">]</span></span></span></em>, <em class="sig-param"><span class="n"><span class="pre">spatial_dims</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">3</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">base_patch_size</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">]</span></span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">[4,</span> <span class="pre">8,</span> <span class="pre">8]</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">enc_dim</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">768</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">emb_dim</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">192</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">num_layer</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">4</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">num_head</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">3</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">learnable_pos_embedding</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">bool</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">True</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/cross_mae.html#CrossMAE_Decoder"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.cross_mae.CrossMAE_Decoder" title="Permalink to this definition">#</a></dt>
-<dd><p>Bases: <code class="xref py py-class docutils literal notranslate"><span class="pre">Module</span></code></p>
-<p>Decoder inspired by [CrossMAE](<a class="reference external" href="https://crossmae.github.io/">https://crossmae.github.io/</a>) where masked tokens only attend
-to visible tokens.</p>
-<dl class="field-list simple">
-<dt class="field-odd">Parameters<span class="colon">:</span></dt>
-<dd class="field-odd"><ul class="simple">
-<li><p><strong>num_patches</strong> (<em>List[int]</em>) – Number of patches in each dimension</p></li>
-<li><p><strong>base_patch_size</strong> (<em>Tuple[int]</em>) – Size of each patch</p></li>
-<li><p><strong>enc_dim</strong> (<em>int</em>) – Dimension of encoder</p></li>
-<li><p><strong>emb_dim</strong> (<em>int</em>) – Dimension of embedding</p></li>
-<li><p><strong>num_layer</strong> (<em>int</em>) – Number of transformer layers</p></li>
-<li><p><strong>num_head</strong> (<em>int</em>) – Number of heads in transformer</p></li>
-<li><p><strong>learnable_pos_embedding</strong> (<em>bool</em>) – If True, learnable positional embeddings are used. If False, fixed sin/cos positional embeddings are used. Empirically, fixed positional embeddings work better for brightfield images.</p></li>
-</ul>
-</dd>
-</dl>
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.blocks.patchify.patchify.Patchify">
+<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">cyto_dl.nn.vits.blocks.patchify.patchify.</span></span><span class="sig-name descname"><span class="pre">Patchify</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">patch_size</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">]</span></span></span></em>, <em class="sig-param"><span class="n"><span class="pre">emb_dim</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">n_patches</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">]</span></span></span></em>, <em class="sig-param"><span class="n"><span class="pre">spatial_dims</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">3</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">context_pixels</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">]</span></span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">[0,</span> <span class="pre">0,</span> <span class="pre">0]</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">input_channels</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">1</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">tasks</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">str</span><span class="p"><span class="pre">]</span></span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">[]</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">learnable_pos_embedding</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">bool</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">True</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/blocks/patchify/patchify.html#Patchify"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.blocks.patchify.patchify.Patchify" title="Permalink to this definition">#</a></dt>
+<dd><p>Bases: <a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html#cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase" title="cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase"><code class="xref py py-class docutils literal notranslate"><span class="pre">PatchifyBase</span></code></a></p>
+<p>Class for converting images to a masked sequence of patches with positional embeddings.</p>
 <dl class="py method">
-<dt class="sig sig-object py" id="cyto_dl.nn.vits.cross_mae.CrossMAE_Decoder.forward">
-<span class="sig-name descname"><span class="pre">forward</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">features</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">forward_indexes</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">backward_indexes</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/cross_mae.html#CrossMAE_Decoder.forward"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.cross_mae.CrossMAE_Decoder.forward" title="Permalink to this definition">#</a></dt>
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.blocks.patchify.patchify.Patchify.create_img2token">
+<span class="sig-name descname"><span class="pre">create_img2token</span></span><span class="sig-paren">(</span><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/blocks/patchify/patchify.html#Patchify.create_img2token"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.blocks.patchify.patchify.Patchify.create_img2token" title="Permalink to this definition">#</a></dt>
+<dd><p>Rearranges the image tensor to a sequence of patches.</p>
+</dd></dl>
+
+<dl class="py method">
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.blocks.patchify.patchify.Patchify.extract_visible_tokens">
+<span class="sig-name descname"><span class="pre">extract_visible_tokens</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">tokens</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">forward_indexes</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">n_visible_patches</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/blocks/patchify/patchify.html#Patchify.extract_visible_tokens"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.blocks.patchify.patchify.Patchify.extract_visible_tokens" title="Permalink to this definition">#</a></dt>
 <dd></dd></dl>
 
 <dl class="py method">
-<dt class="sig sig-object py" id="cyto_dl.nn.vits.cross_mae.CrossMAE_Decoder.init_weight">
-<span class="sig-name descname"><span class="pre">init_weight</span></span><span class="sig-paren">(</span><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/cross_mae.html#CrossMAE_Decoder.init_weight"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.cross_mae.CrossMAE_Decoder.init_weight" title="Permalink to this definition">#</a></dt>
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.blocks.patchify.patchify.Patchify.get_mask_args">
+<span class="sig-name descname"><span class="pre">get_mask_args</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">mask_ratio</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/blocks/patchify/patchify.html#Patchify.get_mask_args"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.blocks.patchify.patchify.Patchify.get_mask_args" title="Permalink to this definition">#</a></dt>
+<dd></dd></dl>
+
+<dl class="py property">
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.blocks.patchify.patchify.Patchify.img2token">
+<em class="property"><span class="pre">property</span><span class="w"> </span></em><span class="sig-name descname"><span class="pre">img2token</span></span><a class="headerlink" href="#cyto_dl.nn.vits.blocks.patchify.patchify.Patchify.img2token" title="Permalink to this definition">#</a></dt>
 <dd></dd></dl>
 
 </dd></dl>
@@ -474,12 +478,12 @@
       <footer>
         
         <div class="related-pages">
-          <a class="next-page" href="cyto_dl.nn.vits.jepa.html">
+          <a class="next-page" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">
               <div class="page-info">
                 <div class="context">
                   <span>Next</span>
                 </div>
-                <div class="title">cyto_dl.nn.vits.jepa module</div>
+                <div class="title">cyto_dl.nn.vits.blocks.patchify.patchify_base module</div>
               </div>
               <svg class="furo-related-icon"><use href="#svg-arrow-right"></use></svg>
             </a>
@@ -490,7 +494,7 @@
                   <span>Previous</span>
                 </div>
                 
-                <div class="title">cyto_dl.nn.vits.blocks.patchify module</div>
+                <div class="title">cyto_dl.nn.vits.blocks.patchify package</div>
                 
               </div>
             </a>
@@ -524,10 +528,12 @@
         <div class="toc-tree-container">
           <div class="toc-tree">
             <ul>
-<li><a class="reference internal" href="#">cyto_dl.nn.vits.cross_mae module</a><ul>
-<li><a class="reference internal" href="#cyto_dl.nn.vits.cross_mae.CrossMAE_Decoder"><code class="docutils literal notranslate"><span class="pre">CrossMAE_Decoder</span></code></a><ul>
-<li><a class="reference internal" href="#cyto_dl.nn.vits.cross_mae.CrossMAE_Decoder.forward"><code class="docutils literal notranslate"><span class="pre">CrossMAE_Decoder.forward()</span></code></a></li>
-<li><a class="reference internal" href="#cyto_dl.nn.vits.cross_mae.CrossMAE_Decoder.init_weight"><code class="docutils literal notranslate"><span class="pre">CrossMAE_Decoder.init_weight()</span></code></a></li>
+<li><a class="reference internal" href="#">cyto_dl.nn.vits.blocks.patchify.patchify module</a><ul>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.blocks.patchify.patchify.Patchify"><code class="docutils literal notranslate"><span class="pre">Patchify</span></code></a><ul>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.blocks.patchify.patchify.Patchify.create_img2token"><code class="docutils literal notranslate"><span class="pre">Patchify.create_img2token()</span></code></a></li>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.blocks.patchify.patchify.Patchify.extract_visible_tokens"><code class="docutils literal notranslate"><span class="pre">Patchify.extract_visible_tokens()</span></code></a></li>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.blocks.patchify.patchify.Patchify.get_mask_args"><code class="docutils literal notranslate"><span class="pre">Patchify.get_mask_args()</span></code></a></li>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.blocks.patchify.patchify.Patchify.img2token"><code class="docutils literal notranslate"><span class="pre">Patchify.img2token</span></code></a></li>
 </ul>
 </li>
 </ul>
diff --git a/cyto_dl.nn.vits.blocks.patchify.patchify_base.html b/cyto_dl.nn.vits.blocks.patchify.patchify_base.html
new file mode 100644
index 000000000..fb6a49135
--- /dev/null
+++ b/cyto_dl.nn.vits.blocks.patchify.patchify_base.html
@@ -0,0 +1,588 @@
+<!doctype html>
+<html class="no-js" lang="en">
+  <head><meta charset="utf-8"/>
+    <meta name="viewport" content="width=device-width,initial-scale=1"/>
+    <meta name="color-scheme" content="light dark"><meta name="generator" content="Docutils 0.19: https://docutils.sourceforge.io/" />
+<link rel="index" title="Index" href="genindex.html" /><link rel="search" title="Search" href="search.html" /><link rel="next" title="cyto_dl.nn.vits.blocks.patchify.patchify_hiera module" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html" /><link rel="prev" title="cyto_dl.nn.vits.blocks.patchify.patchify module" href="cyto_dl.nn.vits.blocks.patchify.patchify.html" />
+
+    <!-- Generated with Sphinx 6.1.3 and Furo 2023.03.23 -->
+        <title>cyto_dl.nn.vits.blocks.patchify.patchify_base module - cyto-dl 0.4.0 documentation</title>
+      <link rel="stylesheet" type="text/css" href="_static/pygments.css" />
+    <link rel="stylesheet" type="text/css" href="_static/styles/furo.css?digest=fad236701ea90a88636c2a8c73b44ae642ed2a53" />
+    <link rel="stylesheet" type="text/css" href="_static/styles/furo-extensions.css?digest=30d1aed668e5c3a91c3e3bf6a60b675221979f0e" />
+    
+    
+
+
+<style>
+  body {
+    --color-code-background: #eeffcc;
+  --color-code-foreground: black;
+  
+  }
+  @media not print {
+    body[data-theme="dark"] {
+      --color-code-background: #202020;
+  --color-code-foreground: #d0d0d0;
+  
+    }
+    @media (prefers-color-scheme: dark) {
+      body:not([data-theme="light"]) {
+        --color-code-background: #202020;
+  --color-code-foreground: #d0d0d0;
+  
+      }
+    }
+  }
+</style></head>
+  <body>
+    
+    <script>
+      document.body.dataset.theme = localStorage.getItem("theme") || "auto";
+    </script>
+    
+
+<svg xmlns="http://www.w3.org/2000/svg" style="display: none;">
+  <symbol id="svg-toc" viewBox="0 0 24 24">
+    <title>Contents</title>
+    <svg stroke="currentColor" fill="currentColor" stroke-width="0" viewBox="0 0 1024 1024">
+      <path d="M408 442h480c4.4 0 8-3.6 8-8v-56c0-4.4-3.6-8-8-8H408c-4.4 0-8 3.6-8 8v56c0 4.4 3.6 8 8 8zm-8 204c0 4.4 3.6 8 8 8h480c4.4 0 8-3.6 8-8v-56c0-4.4-3.6-8-8-8H408c-4.4 0-8 3.6-8 8v56zm504-486H120c-4.4 0-8 3.6-8 8v56c0 4.4 3.6 8 8 8h784c4.4 0 8-3.6 8-8v-56c0-4.4-3.6-8-8-8zm0 632H120c-4.4 0-8 3.6-8 8v56c0 4.4 3.6 8 8 8h784c4.4 0 8-3.6 8-8v-56c0-4.4-3.6-8-8-8zM115.4 518.9L271.7 642c5.8 4.6 14.4.5 14.4-6.9V388.9c0-7.4-8.5-11.5-14.4-6.9L115.4 505.1a8.74 8.74 0 0 0 0 13.8z"/>
+    </svg>
+  </symbol>
+  <symbol id="svg-menu" viewBox="0 0 24 24">
+    <title>Menu</title>
+    <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
+      stroke-width="2" stroke-linecap="round" stroke-linejoin="round" class="feather-menu">
+      <line x1="3" y1="12" x2="21" y2="12"></line>
+      <line x1="3" y1="6" x2="21" y2="6"></line>
+      <line x1="3" y1="18" x2="21" y2="18"></line>
+    </svg>
+  </symbol>
+  <symbol id="svg-arrow-right" viewBox="0 0 24 24">
+    <title>Expand</title>
+    <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
+      stroke-width="2" stroke-linecap="round" stroke-linejoin="round" class="feather-chevron-right">
+      <polyline points="9 18 15 12 9 6"></polyline>
+    </svg>
+  </symbol>
+  <symbol id="svg-sun" viewBox="0 0 24 24">
+    <title>Light mode</title>
+    <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
+      stroke-width="1.5" stroke-linecap="round" stroke-linejoin="round" class="feather-sun">
+      <circle cx="12" cy="12" r="5"></circle>
+      <line x1="12" y1="1" x2="12" y2="3"></line>
+      <line x1="12" y1="21" x2="12" y2="23"></line>
+      <line x1="4.22" y1="4.22" x2="5.64" y2="5.64"></line>
+      <line x1="18.36" y1="18.36" x2="19.78" y2="19.78"></line>
+      <line x1="1" y1="12" x2="3" y2="12"></line>
+      <line x1="21" y1="12" x2="23" y2="12"></line>
+      <line x1="4.22" y1="19.78" x2="5.64" y2="18.36"></line>
+      <line x1="18.36" y1="5.64" x2="19.78" y2="4.22"></line>
+    </svg>
+  </symbol>
+  <symbol id="svg-moon" viewBox="0 0 24 24">
+    <title>Dark mode</title>
+    <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
+      stroke-width="1.5" stroke-linecap="round" stroke-linejoin="round" class="icon-tabler-moon">
+      <path stroke="none" d="M0 0h24v24H0z" fill="none" />
+      <path d="M12 3c.132 0 .263 0 .393 0a7.5 7.5 0 0 0 7.92 12.446a9 9 0 1 1 -8.313 -12.454z" />
+    </svg>
+  </symbol>
+  <symbol id="svg-sun-half" viewBox="0 0 24 24">
+    <title>Auto light/dark mode</title>
+    <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
+      stroke-width="1.5" stroke-linecap="round" stroke-linejoin="round" class="icon-tabler-shadow">
+      <path stroke="none" d="M0 0h24v24H0z" fill="none"/>
+      <circle cx="12" cy="12" r="9" />
+      <path d="M13 12h5" />
+      <path d="M13 15h4" />
+      <path d="M13 18h1" />
+      <path d="M13 9h4" />
+      <path d="M13 6h1" />
+    </svg>
+  </symbol>
+</svg>
+
+<input type="checkbox" class="sidebar-toggle" name="__navigation" id="__navigation">
+<input type="checkbox" class="sidebar-toggle" name="__toc" id="__toc">
+<label class="overlay sidebar-overlay" for="__navigation">
+  <div class="visually-hidden">Hide navigation sidebar</div>
+</label>
+<label class="overlay toc-overlay" for="__toc">
+  <div class="visually-hidden">Hide table of contents sidebar</div>
+</label>
+
+
+
+<div class="page">
+  <header class="mobile-header">
+    <div class="header-left">
+      <label class="nav-overlay-icon" for="__navigation">
+        <div class="visually-hidden">Toggle site navigation sidebar</div>
+        <i class="icon"><svg><use href="#svg-menu"></use></svg></i>
+      </label>
+    </div>
+    <div class="header-center">
+      <a href="index.html"><div class="brand">cyto-dl 0.4.0 documentation</div></a>
+    </div>
+    <div class="header-right">
+      <div class="theme-toggle-container theme-toggle-header">
+        <button class="theme-toggle">
+          <div class="visually-hidden">Toggle Light / Dark / Auto color theme</div>
+          <svg class="theme-icon-when-auto"><use href="#svg-sun-half"></use></svg>
+          <svg class="theme-icon-when-dark"><use href="#svg-moon"></use></svg>
+          <svg class="theme-icon-when-light"><use href="#svg-sun"></use></svg>
+        </button>
+      </div>
+      <label class="toc-overlay-icon toc-header-icon" for="__toc">
+        <div class="visually-hidden">Toggle table of contents sidebar</div>
+        <i class="icon"><svg><use href="#svg-toc"></use></svg></i>
+      </label>
+    </div>
+  </header>
+  <aside class="sidebar-drawer">
+    <div class="sidebar-container">
+      
+      <div class="sidebar-sticky"><a class="sidebar-brand" href="index.html">
+  
+  
+  <span class="sidebar-brand-text">cyto-dl 0.4.0 documentation</span>
+  
+</a><form class="sidebar-search-container" method="get" action="search.html" role="search">
+  <input class="sidebar-search" placeholder="Search" name="q" aria-label="Search">
+  <input type="hidden" name="check_keywords" value="yes">
+  <input type="hidden" name="area" value="default">
+</form>
+<div id="searchbox"></div><div class="sidebar-scroll"><div class="sidebar-tree">
+  <p class="caption" role="heading"><span class="caption-text">Contents:</span></p>
+<ul class="current">
+<li class="toctree-l1"><a class="reference internal" href="index.html">Home</a></li>
+<li class="toctree-l1"><a class="reference internal" href="installation.html">Installation</a></li>
+<li class="toctree-l1"><a class="reference internal" href="using_examples.html">Modifying the Example Configs</a></li>
+<li class="toctree-l1"><a class="reference internal" href="api.html">Using the API</a></li>
+<li class="toctree-l1 current has-children"><a class="reference internal" href="modules.html">Package modules</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-1" name="toctree-checkbox-1" role="switch" type="checkbox"/><label for="toctree-checkbox-1"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
+<li class="toctree-l2 current has-children"><a class="reference internal" href="cyto_dl.html">cyto_dl package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-2" name="toctree-checkbox-2" role="switch" type="checkbox"/><label for="toctree-checkbox-2"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.api.html">cyto_dl.api package</a><input class="toctree-checkbox" id="toctree-checkbox-3" name="toctree-checkbox-3" role="switch" type="checkbox"/><label for="toctree-checkbox-3"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.api.cyto_dl_model.html">cyto_dl.api.cyto_dl_model package</a><input class="toctree-checkbox" id="toctree-checkbox-4" name="toctree-checkbox-4" role="switch" type="checkbox"/><label for="toctree-checkbox-4"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.api.cyto_dl_model.cyto_dl_base_model.html">cyto_dl.api.cyto_dl_model.cyto_dl_base_model module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.api.cyto_dl_model.segmentation_plugin_model.html">cyto_dl.api.cyto_dl_model.segmentation_plugin_model module</a></li>
+</ul>
+</li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.api.data.html">cyto_dl.api.data module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.api.model.html">cyto_dl.api.model module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.callbacks.html">cyto_dl.callbacks package</a><input class="toctree-checkbox" id="toctree-checkbox-5" name="toctree-checkbox-5" role="switch" type="checkbox"/><label for="toctree-checkbox-5"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.callbacks.callback_utils.html">cyto_dl.callbacks.callback_utils module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.callbacks.csv_saver.html">cyto_dl.callbacks.csv_saver module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.callbacks.image_saver.html">cyto_dl.callbacks.image_saver module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.callbacks.latent_walk.html">cyto_dl.callbacks.latent_walk module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.callbacks.layer_freeze.html">cyto_dl.callbacks.layer_freeze module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.callbacks.model_utils.html">cyto_dl.callbacks.model_utils module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.callbacks.outlier_detection.html">cyto_dl.callbacks.outlier_detection module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.dataframe.html">cyto_dl.dataframe package</a><input class="toctree-checkbox" id="toctree-checkbox-6" name="toctree-checkbox-6" role="switch" type="checkbox"/><label for="toctree-checkbox-6"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.dataframe.transforms.html">cyto_dl.dataframe.transforms package</a><input class="toctree-checkbox" id="toctree-checkbox-7" name="toctree-checkbox-7" role="switch" type="checkbox"/><label for="toctree-checkbox-7"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.dataframe.transforms.filter.html">cyto_dl.dataframe.transforms.filter module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.dataframe.transforms.group_cols.html">cyto_dl.dataframe.transforms.group_cols module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.dataframe.transforms.misc.html">cyto_dl.dataframe.transforms.misc module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.dataframe.transforms.split.html">cyto_dl.dataframe.transforms.split module</a></li>
+</ul>
+</li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.dataframe.readers.html">cyto_dl.dataframe.readers module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.datamodules.html">cyto_dl.datamodules package</a><input class="toctree-checkbox" id="toctree-checkbox-8" name="toctree-checkbox-8" role="switch" type="checkbox"/><label for="toctree-checkbox-8"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.datamodules.dataframe.html">cyto_dl.datamodules.dataframe package</a><input class="toctree-checkbox" id="toctree-checkbox-9" name="toctree-checkbox-9" role="switch" type="checkbox"/><label for="toctree-checkbox-9"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.datamodules.dataframe.dataframe_datamodule.html">cyto_dl.datamodules.dataframe.dataframe_datamodule module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.datamodules.dataframe.grouped_dataframe_datamodule.html">cyto_dl.datamodules.dataframe.grouped_dataframe_datamodule module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.datamodules.dataframe.utils.html">cyto_dl.datamodules.dataframe.utils module</a></li>
+</ul>
+</li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.datamodules.array.html">cyto_dl.datamodules.array module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.datamodules.data_dict.html">cyto_dl.datamodules.data_dict module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.datamodules.folder.html">cyto_dl.datamodules.folder module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.datamodules.multidim_image.html">cyto_dl.datamodules.multidim_image module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.datamodules.smartcache.html">cyto_dl.datamodules.smartcache module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.datamodules.torchvision.html">cyto_dl.datamodules.torchvision module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.image.html">cyto_dl.image package</a><input class="toctree-checkbox" id="toctree-checkbox-10" name="toctree-checkbox-10" role="switch" type="checkbox"/><label for="toctree-checkbox-10"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.image.io.html">cyto_dl.image.io package</a><input class="toctree-checkbox" id="toctree-checkbox-11" name="toctree-checkbox-11" role="switch" type="checkbox"/><label for="toctree-checkbox-11"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.io.aicsimage_loader.html">cyto_dl.image.io.aicsimage_loader module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.io.monai_bio_reader.html">cyto_dl.image.io.monai_bio_reader module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.io.numpy_reader.html">cyto_dl.image.io.numpy_reader module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.io.ome_zarr_reader.html">cyto_dl.image.io.ome_zarr_reader module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.io.polygon_loader.html">cyto_dl.image.io.polygon_loader module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.io.skimage_reader.html">cyto_dl.image.io.skimage_reader module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.image.transforms.html">cyto_dl.image.transforms package</a><input class="toctree-checkbox" id="toctree-checkbox-12" name="toctree-checkbox-12" role="switch" type="checkbox"/><label for="toctree-checkbox-12"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.add_meta.html">cyto_dl.image.transforms.add_meta module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.bright_sampler.html">cyto_dl.image.transforms.bright_sampler module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.clip.html">cyto_dl.image.transforms.clip module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.contrastadjust.html">cyto_dl.image.transforms.contrastadjust module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.generate_jepa_masks.html">cyto_dl.image.transforms.generate_jepa_masks module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.merge.html">cyto_dl.image.transforms.merge module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.multiscale_cropper.html">cyto_dl.image.transforms.multiscale_cropper module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.pad.html">cyto_dl.image.transforms.pad module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.project.html">cyto_dl.image.transforms.project module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.rotation_mask_transform.html">cyto_dl.image.transforms.rotation_mask_transform module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.save.html">cyto_dl.image.transforms.save module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.so2_random_rotation.html">cyto_dl.image.transforms.so2_random_rotation module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.track_transforms.html">cyto_dl.image.transforms.track_transforms module</a></li>
+</ul>
+</li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.loggers.html">cyto_dl.loggers package</a><input class="toctree-checkbox" id="toctree-checkbox-13" name="toctree-checkbox-13" role="switch" type="checkbox"/><label for="toctree-checkbox-13"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.loggers.mlflow.html">cyto_dl.loggers.mlflow module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.models.html">cyto_dl.models package</a><input class="toctree-checkbox" id="toctree-checkbox-14" name="toctree-checkbox-14" role="switch" type="checkbox"/><label for="toctree-checkbox-14"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.models.classification.html">cyto_dl.models.classification package</a><input class="toctree-checkbox" id="toctree-checkbox-15" name="toctree-checkbox-15" role="switch" type="checkbox"/><label for="toctree-checkbox-15"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.classification.classification.html">cyto_dl.models.classification.classification module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.classification.timepoint_classification.html">cyto_dl.models.classification.timepoint_classification module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.models.contrastive.html">cyto_dl.models.contrastive package</a><input class="toctree-checkbox" id="toctree-checkbox-16" name="toctree-checkbox-16" role="switch" type="checkbox"/><label for="toctree-checkbox-16"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.contrastive.contrastive.html">cyto_dl.models.contrastive.contrastive module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.models.handlers.html">cyto_dl.models.handlers package</a><input class="toctree-checkbox" id="toctree-checkbox-17" name="toctree-checkbox-17" role="switch" type="checkbox"/><label for="toctree-checkbox-17"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.handlers.base_handler.html">cyto_dl.models.handlers.base_handler module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.handlers.image_handler.html">cyto_dl.models.handlers.image_handler module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.handlers.load_image_patch.html">cyto_dl.models.handlers.load_image_patch module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.models.im2im.html">cyto_dl.models.im2im package</a><input class="toctree-checkbox" id="toctree-checkbox-18" name="toctree-checkbox-18" role="switch" type="checkbox"/><label for="toctree-checkbox-18"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.models.im2im.utils.html">cyto_dl.models.im2im.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-19" name="toctree-checkbox-19" role="switch" type="checkbox"/><label for="toctree-checkbox-19"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.models.im2im.utils.postprocessing.html">cyto_dl.models.im2im.utils.postprocessing package</a><input class="toctree-checkbox" id="toctree-checkbox-20" name="toctree-checkbox-20" role="switch" type="checkbox"/><label for="toctree-checkbox-20"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.models.im2im.utils.postprocessing.act_thresh_label.html">cyto_dl.models.im2im.utils.postprocessing.act_thresh_label module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.models.im2im.utils.postprocessing.auto_thresh.html">cyto_dl.models.im2im.utils.postprocessing.auto_thresh module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.models.im2im.utils.postprocessing.dict_to_im.html">cyto_dl.models.im2im.utils.postprocessing.dict_to_im module</a></li>
+</ul>
+</li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.models.im2im.utils.instance_seg.html">cyto_dl.models.im2im.utils.instance_seg module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.models.im2im.utils.noise_annealer.html">cyto_dl.models.im2im.utils.noise_annealer module</a></li>
+</ul>
+</li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.im2im.gan.html">cyto_dl.models.im2im.gan module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.im2im.multi_task.html">cyto_dl.models.im2im.multi_task module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.models.jepa.html">cyto_dl.models.jepa package</a><input class="toctree-checkbox" id="toctree-checkbox-21" name="toctree-checkbox-21" role="switch" type="checkbox"/><label for="toctree-checkbox-21"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.jepa.ijepa.html">cyto_dl.models.jepa.ijepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.jepa.iwm.html">cyto_dl.models.jepa.iwm module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.jepa.jepa_base.html">cyto_dl.models.jepa.jepa_base module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.models.utils.html">cyto_dl.models.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-22" name="toctree-checkbox-22" role="switch" type="checkbox"/><label for="toctree-checkbox-22"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.utils.mlflow.html">cyto_dl.models.utils.mlflow module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.utils.utils.html">cyto_dl.models.utils.utils module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.models.vae.html">cyto_dl.models.vae package</a><input class="toctree-checkbox" id="toctree-checkbox-23" name="toctree-checkbox-23" role="switch" type="checkbox"/><label for="toctree-checkbox-23"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.models.vae.o2_spharm_vae.html">cyto_dl.models.vae.o2_spharm_vae package</a><input class="toctree-checkbox" id="toctree-checkbox-24" name="toctree-checkbox-24" role="switch" type="checkbox"/><label for="toctree-checkbox-24"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.models.vae.o2_spharm_vae.o2_spharm_encoder.html">cyto_dl.models.vae.o2_spharm_vae.o2_spharm_encoder module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.models.vae.o2_spharm_vae.o2_spharm_vae.html">cyto_dl.models.vae.o2_spharm_vae.o2_spharm_vae module</a></li>
+</ul>
+</li>
+<li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.models.vae.priors.html">cyto_dl.models.vae.priors package</a><input class="toctree-checkbox" id="toctree-checkbox-25" name="toctree-checkbox-25" role="switch" type="checkbox"/><label for="toctree-checkbox-25"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.models.vae.priors.abstract_prior.html">cyto_dl.models.vae.priors.abstract_prior module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.models.vae.priors.gaussian.html">cyto_dl.models.vae.priors.gaussian module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.models.vae.priors.identity_prior.html">cyto_dl.models.vae.priors.identity_prior module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.models.vae.priors.joint_prior.html">cyto_dl.models.vae.priors.joint_prior module</a></li>
+</ul>
+</li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.vae.base_vae.html">cyto_dl.models.vae.base_vae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.vae.image_canon_vae.html">cyto_dl.models.vae.image_canon_vae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.vae.image_encoder.html">cyto_dl.models.vae.image_encoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.vae.image_vae.html">cyto_dl.models.vae.image_vae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.vae.latent_loss_vae.html">cyto_dl.models.vae.latent_loss_vae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.vae.point_cloud_vae.html">cyto_dl.models.vae.point_cloud_vae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.vae.tabular_vae.html">cyto_dl.models.vae.tabular_vae module</a></li>
+</ul>
+</li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.models.base_model.html">cyto_dl.models.base_model module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.models.basic_model.html">cyto_dl.models.basic_model module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 current has-children"><a class="reference internal" href="cyto_dl.nn.html">cyto_dl.nn package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-26" name="toctree-checkbox-26" role="switch" type="checkbox"/><label for="toctree-checkbox-26"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.discriminators.html">cyto_dl.nn.discriminators package</a><input class="toctree-checkbox" id="toctree-checkbox-27" name="toctree-checkbox-27" role="switch" type="checkbox"/><label for="toctree-checkbox-27"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.discriminators.multi_scale_discriminator.html">cyto_dl.nn.discriminators.multi_scale_discriminator module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.discriminators.n_layer_discriminator.html">cyto_dl.nn.discriminators.n_layer_discriminator module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.head.html">cyto_dl.nn.head package</a><input class="toctree-checkbox" id="toctree-checkbox-28" name="toctree-checkbox-28" role="switch" type="checkbox"/><label for="toctree-checkbox-28"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.head.base_head.html">cyto_dl.nn.head.base_head module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.head.gan_head.html">cyto_dl.nn.head.gan_head module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.head.gan_head_superres.html">cyto_dl.nn.head.gan_head_superres module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.head.mae_head.html">cyto_dl.nn.head.mae_head module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.head.mask_head.html">cyto_dl.nn.head.mask_head module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.head.res_blocks_head.html">cyto_dl.nn.head.res_blocks_head module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.head.vic_reg.html">cyto_dl.nn.head.vic_reg module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.losses.html">cyto_dl.nn.losses package</a><input class="toctree-checkbox" id="toctree-checkbox-29" name="toctree-checkbox-29" role="switch" type="checkbox"/><label for="toctree-checkbox-29"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.adversarial_loss.html">cyto_dl.nn.losses.adversarial_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.chamfer_loss.html">cyto_dl.nn.losses.chamfer_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.continuous_bernoulli.html">cyto_dl.nn.losses.continuous_bernoulli module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.cosine_loss.html">cyto_dl.nn.losses.cosine_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.gan_loss.html">cyto_dl.nn.losses.gan_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.gaussian_nll_loss.html">cyto_dl.nn.losses.gaussian_nll_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.geomloss.html">cyto_dl.nn.losses.geomloss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.loss_wrapper.html">cyto_dl.nn.losses.loss_wrapper module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.spharm_loss.html">cyto_dl.nn.losses.spharm_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.threshold_loss.html">cyto_dl.nn.losses.threshold_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.vic_reg.html">cyto_dl.nn.losses.vic_reg module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.weibull.html">cyto_dl.nn.losses.weibull module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.weighted_mse_loss.html">cyto_dl.nn.losses.weighted_mse_loss module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.point_cloud.html">cyto_dl.nn.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-30" name="toctree-checkbox-30" role="switch" type="checkbox"/><label for="toctree-checkbox-30"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.point_cloud.dgcnn.html">cyto_dl.nn.point_cloud.dgcnn module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.point_cloud.folding_net.html">cyto_dl.nn.point_cloud.folding_net module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.point_cloud.graph_functions.html">cyto_dl.nn.point_cloud.graph_functions module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.point_cloud.vnn.html">cyto_dl.nn.point_cloud.vnn module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 current has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
+<li class="toctree-l5 current has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
+<li class="toctree-l6 current has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7 current current-page"><a class="current reference internal" href="#">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
+</ul>
+</li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
+</ul>
+</li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.hr_skip.html">cyto_dl.nn.hr_skip module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.mlp.html">cyto_dl.nn.mlp module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.res_unit.html">cyto_dl.nn.res_unit module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.spatial_transformer.html">cyto_dl.nn.spatial_transformer module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
+</ul>
+</li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
+</ul>
+</li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.utils.array.html">cyto_dl.utils.array module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.utils.checkpoint.html">cyto_dl.utils.checkpoint module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.utils.config.html">cyto_dl.utils.config module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.utils.download_test_data.html">cyto_dl.utils.download_test_data module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.utils.dummy_dataset.html">cyto_dl.utils.dummy_dataset module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.utils.pylogger.html">cyto_dl.utils.pylogger module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.utils.rich_utils.html">cyto_dl.utils.rich_utils module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.utils.rotation.html">cyto_dl.utils.rotation module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.utils.template_utils.html">cyto_dl.utils.template_utils module</a></li>
+</ul>
+</li>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.compile.html">cyto_dl.compile module</a></li>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.eval.html">cyto_dl.eval module</a></li>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.train.html">cyto_dl.train module</a></li>
+</ul>
+</li>
+</ul>
+</li>
+</ul>
+
+</div>
+</div>
+
+      </div>
+      
+    </div>
+  </aside>
+  <div class="main">
+    <div class="content">
+      <div class="article-container">
+        <a href="#" class="back-to-top muted-link">
+          <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24">
+            <path d="M13 20h-2V8l-5.5 5.5-1.42-1.42L12 4.16l7.92 7.92-1.42 1.42L13 8v12z"></path>
+          </svg>
+          <span>Back to top</span>
+        </a>
+        <div class="content-icon-container">
+          
+<div class="theme-toggle-container theme-toggle-content">
+            <button class="theme-toggle">
+              <div class="visually-hidden">Toggle Light / Dark / Auto color theme</div>
+              <svg class="theme-icon-when-auto"><use href="#svg-sun-half"></use></svg>
+              <svg class="theme-icon-when-dark"><use href="#svg-moon"></use></svg>
+              <svg class="theme-icon-when-light"><use href="#svg-sun"></use></svg>
+            </button>
+          </div>
+          <label class="toc-overlay-icon toc-content-icon" for="__toc">
+            <div class="visually-hidden">Toggle table of contents sidebar</div>
+            <i class="icon"><svg><use href="#svg-toc"></use></svg></i>
+          </label>
+        </div>
+        <article role="main">
+          <section id="module-cyto_dl.nn.vits.blocks.patchify.patchify_base">
+<span id="cyto-dl-nn-vits-blocks-patchify-patchify-base-module"></span><h1>cyto_dl.nn.vits.blocks.patchify.patchify_base module<a class="headerlink" href="#module-cyto_dl.nn.vits.blocks.patchify.patchify_base" title="Permalink to this heading">#</a></h1>
+<dl class="py class">
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase">
+<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">cyto_dl.nn.vits.blocks.patchify.patchify_base.</span></span><span class="sig-name descname"><span class="pre">PatchifyBase</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">patch_size</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">]</span></span></span></em>, <em class="sig-param"><span class="n"><span class="pre">emb_dim</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">n_patches</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">]</span></span></span></em>, <em class="sig-param"><span class="n"><span class="pre">spatial_dims</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">3</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">context_pixels</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">]</span></span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">[0,</span> <span class="pre">0,</span> <span class="pre">0]</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">input_channels</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">1</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">tasks</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">str</span><span class="p"><span class="pre">]</span></span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">[]</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">learnable_pos_embedding</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">bool</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">True</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/blocks/patchify/patchify_base.html#PatchifyBase"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase" title="Permalink to this definition">#</a></dt>
+<dd><p>Bases: <code class="xref py py-class docutils literal notranslate"><span class="pre">Module</span></code>, <code class="xref py py-class docutils literal notranslate"><span class="pre">ABC</span></code></p>
+<p>Class for converting images to a masked sequence of patches with positional embeddings.</p>
+<dl class="field-list simple">
+<dt class="field-odd">Parameters<span class="colon">:</span></dt>
+<dd class="field-odd"><ul class="simple">
+<li><p><strong>patch_size</strong> (<em>List[int]</em>) – Size of each patch in pix (ZYX order for 3D, YX order for 2D)</p></li>
+<li><p><strong>emb_dim</strong> (<em>int</em>) – Dimension of encoder</p></li>
+<li><p><strong>n_patches</strong> (<em>List[int]</em>) – Number of patches in each spatial dimension (ZYX order for 3D, YX order for 2D)</p></li>
+<li><p><strong>spatial_dims</strong> (<em>int</em>) – Number of spatial dimensions</p></li>
+<li><p><strong>context_pixels</strong> (<em>List[int]</em>) – Number of extra pixels around each patch to include in convolutional embedding to encoder dimension.</p></li>
+<li><p><strong>input_channels</strong> (<em>int</em>) – Number of input channels</p></li>
+<li><p><strong>tasks</strong> (<em>List[str]</em>) – List of tasks to encode</p></li>
+<li><p><strong>learnable_pos_embedding</strong> (<em>bool</em>) – If True, learnable positional embeddings are used. If False, fixed sin/cos positional embeddings. Empirically, fixed positional embeddings work better for brightfield images.</p></li>
+</ul>
+</dd>
+</dl>
+<dl class="py method">
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase.create_conv">
+<span class="sig-name descname"><span class="pre">create_conv</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">input_channels</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">emb_dim</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">patch_size</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">context_pixels</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/blocks/patchify/patchify_base.html#PatchifyBase.create_conv"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase.create_conv" title="Permalink to this definition">#</a></dt>
+<dd></dd></dl>
+
+<dl class="py method">
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase.create_patch2img">
+<span class="sig-name descname"><span class="pre">create_patch2img</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">n_patches</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">patch_size</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/blocks/patchify/patchify_base.html#PatchifyBase.create_patch2img"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase.create_patch2img" title="Permalink to this definition">#</a></dt>
+<dd><p>Converts boolean array of whether to keep index of each patch to an image-shaped mask of
+same size as input image.</p>
+</dd></dl>
+
+<dl class="py method">
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase.extract_visible_tokens">
+<em class="property"><span class="pre">abstract</span><span class="w"> </span></em><span class="sig-name descname"><span class="pre">extract_visible_tokens</span></span><span class="sig-paren">(</span><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/blocks/patchify/patchify_base.html#PatchifyBase.extract_visible_tokens"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase.extract_visible_tokens" title="Permalink to this definition">#</a></dt>
+<dd></dd></dl>
+
+<dl class="py method">
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase.forward">
+<span class="sig-name descname"><span class="pre">forward</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">img</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">mask_ratio</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">task</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/blocks/patchify/patchify_base.html#PatchifyBase.forward"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase.forward" title="Permalink to this definition">#</a></dt>
+<dd></dd></dl>
+
+<dl class="py method">
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase.get_mask">
+<span class="sig-name descname"><span class="pre">get_mask</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">img</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">n_visible_patches</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">num_patches</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/blocks/patchify/patchify_base.html#PatchifyBase.get_mask"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase.get_mask" title="Permalink to this definition">#</a></dt>
+<dd></dd></dl>
+
+<dl class="py method">
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase.get_mask_args">
+<em class="property"><span class="pre">abstract</span><span class="w"> </span></em><span class="sig-name descname"><span class="pre">get_mask_args</span></span><span class="sig-paren">(</span><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/blocks/patchify/patchify_base.html#PatchifyBase.get_mask_args"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase.get_mask_args" title="Permalink to this definition">#</a></dt>
+<dd></dd></dl>
+
+<dl class="py property">
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase.img2token">
+<em class="property"><span class="pre">abstract</span><span class="w"> </span><span class="pre">property</span><span class="w"> </span></em><span class="sig-name descname"><span class="pre">img2token</span></span><a class="headerlink" href="#cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase.img2token" title="Permalink to this definition">#</a></dt>
+<dd></dd></dl>
+
+</dd></dl>
+
+</section>
+
+        </article>
+      </div>
+      <footer>
+        
+        <div class="related-pages">
+          <a class="next-page" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">
+              <div class="page-info">
+                <div class="context">
+                  <span>Next</span>
+                </div>
+                <div class="title">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</div>
+              </div>
+              <svg class="furo-related-icon"><use href="#svg-arrow-right"></use></svg>
+            </a>
+          <a class="prev-page" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">
+              <svg class="furo-related-icon"><use href="#svg-arrow-right"></use></svg>
+              <div class="page-info">
+                <div class="context">
+                  <span>Previous</span>
+                </div>
+                
+                <div class="title">cyto_dl.nn.vits.blocks.patchify.patchify module</div>
+                
+              </div>
+            </a>
+        </div>
+        <div class="bottom-of-page">
+          <div class="left-details">
+            <div class="copyright">
+                Copyright &#169; 2023
+            </div>
+            Made with <a href="https://www.sphinx-doc.org/">Sphinx</a> and <a class="muted-link" href="https://pradyunsg.me">@pradyunsg</a>'s
+            
+            <a href="https://github.com/pradyunsg/furo">Furo</a>
+            
+          </div>
+          <div class="right-details">
+            
+          </div>
+        </div>
+        
+      </footer>
+    </div>
+    <aside class="toc-drawer">
+      
+      
+      <div class="toc-sticky toc-scroll">
+        <div class="toc-title-container">
+          <span class="toc-title">
+            On this page
+          </span>
+        </div>
+        <div class="toc-tree-container">
+          <div class="toc-tree">
+            <ul>
+<li><a class="reference internal" href="#">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a><ul>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase"><code class="docutils literal notranslate"><span class="pre">PatchifyBase</span></code></a><ul>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase.create_conv"><code class="docutils literal notranslate"><span class="pre">PatchifyBase.create_conv()</span></code></a></li>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase.create_patch2img"><code class="docutils literal notranslate"><span class="pre">PatchifyBase.create_patch2img()</span></code></a></li>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase.extract_visible_tokens"><code class="docutils literal notranslate"><span class="pre">PatchifyBase.extract_visible_tokens()</span></code></a></li>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase.forward"><code class="docutils literal notranslate"><span class="pre">PatchifyBase.forward()</span></code></a></li>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase.get_mask"><code class="docutils literal notranslate"><span class="pre">PatchifyBase.get_mask()</span></code></a></li>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase.get_mask_args"><code class="docutils literal notranslate"><span class="pre">PatchifyBase.get_mask_args()</span></code></a></li>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase.img2token"><code class="docutils literal notranslate"><span class="pre">PatchifyBase.img2token</span></code></a></li>
+</ul>
+</li>
+</ul>
+</li>
+</ul>
+
+          </div>
+        </div>
+      </div>
+      
+      
+    </aside>
+  </div>
+</div><script data-url_root="./" id="documentation_options" src="_static/documentation_options.js"></script>
+    <script src="_static/doctools.js"></script>
+    <script src="_static/sphinx_highlight.js"></script>
+    <script src="_static/scripts/furo.js"></script>
+    </body>
+</html>
\ No newline at end of file
diff --git a/cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html b/cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html
new file mode 100644
index 000000000..932a647f3
--- /dev/null
+++ b/cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html
@@ -0,0 +1,581 @@
+<!doctype html>
+<html class="no-js" lang="en">
+  <head><meta charset="utf-8"/>
+    <meta name="viewport" content="width=device-width,initial-scale=1"/>
+    <meta name="color-scheme" content="light dark"><meta name="generator" content="Docutils 0.19: https://docutils.sourceforge.io/" />
+<link rel="index" title="Index" href="genindex.html" /><link rel="search" title="Search" href="search.html" /><link rel="next" title="cyto_dl.nn.vits.blocks.cross_attention module" href="cyto_dl.nn.vits.blocks.cross_attention.html" /><link rel="prev" title="cyto_dl.nn.vits.blocks.patchify.patchify_base module" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html" />
+
+    <!-- Generated with Sphinx 6.1.3 and Furo 2023.03.23 -->
+        <title>cyto_dl.nn.vits.blocks.patchify.patchify_hiera module - cyto-dl 0.4.0 documentation</title>
+      <link rel="stylesheet" type="text/css" href="_static/pygments.css" />
+    <link rel="stylesheet" type="text/css" href="_static/styles/furo.css?digest=fad236701ea90a88636c2a8c73b44ae642ed2a53" />
+    <link rel="stylesheet" type="text/css" href="_static/styles/furo-extensions.css?digest=30d1aed668e5c3a91c3e3bf6a60b675221979f0e" />
+    
+    
+
+
+<style>
+  body {
+    --color-code-background: #eeffcc;
+  --color-code-foreground: black;
+  
+  }
+  @media not print {
+    body[data-theme="dark"] {
+      --color-code-background: #202020;
+  --color-code-foreground: #d0d0d0;
+  
+    }
+    @media (prefers-color-scheme: dark) {
+      body:not([data-theme="light"]) {
+        --color-code-background: #202020;
+  --color-code-foreground: #d0d0d0;
+  
+      }
+    }
+  }
+</style></head>
+  <body>
+    
+    <script>
+      document.body.dataset.theme = localStorage.getItem("theme") || "auto";
+    </script>
+    
+
+<svg xmlns="http://www.w3.org/2000/svg" style="display: none;">
+  <symbol id="svg-toc" viewBox="0 0 24 24">
+    <title>Contents</title>
+    <svg stroke="currentColor" fill="currentColor" stroke-width="0" viewBox="0 0 1024 1024">
+      <path d="M408 442h480c4.4 0 8-3.6 8-8v-56c0-4.4-3.6-8-8-8H408c-4.4 0-8 3.6-8 8v56c0 4.4 3.6 8 8 8zm-8 204c0 4.4 3.6 8 8 8h480c4.4 0 8-3.6 8-8v-56c0-4.4-3.6-8-8-8H408c-4.4 0-8 3.6-8 8v56zm504-486H120c-4.4 0-8 3.6-8 8v56c0 4.4 3.6 8 8 8h784c4.4 0 8-3.6 8-8v-56c0-4.4-3.6-8-8-8zm0 632H120c-4.4 0-8 3.6-8 8v56c0 4.4 3.6 8 8 8h784c4.4 0 8-3.6 8-8v-56c0-4.4-3.6-8-8-8zM115.4 518.9L271.7 642c5.8 4.6 14.4.5 14.4-6.9V388.9c0-7.4-8.5-11.5-14.4-6.9L115.4 505.1a8.74 8.74 0 0 0 0 13.8z"/>
+    </svg>
+  </symbol>
+  <symbol id="svg-menu" viewBox="0 0 24 24">
+    <title>Menu</title>
+    <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
+      stroke-width="2" stroke-linecap="round" stroke-linejoin="round" class="feather-menu">
+      <line x1="3" y1="12" x2="21" y2="12"></line>
+      <line x1="3" y1="6" x2="21" y2="6"></line>
+      <line x1="3" y1="18" x2="21" y2="18"></line>
+    </svg>
+  </symbol>
+  <symbol id="svg-arrow-right" viewBox="0 0 24 24">
+    <title>Expand</title>
+    <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
+      stroke-width="2" stroke-linecap="round" stroke-linejoin="round" class="feather-chevron-right">
+      <polyline points="9 18 15 12 9 6"></polyline>
+    </svg>
+  </symbol>
+  <symbol id="svg-sun" viewBox="0 0 24 24">
+    <title>Light mode</title>
+    <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
+      stroke-width="1.5" stroke-linecap="round" stroke-linejoin="round" class="feather-sun">
+      <circle cx="12" cy="12" r="5"></circle>
+      <line x1="12" y1="1" x2="12" y2="3"></line>
+      <line x1="12" y1="21" x2="12" y2="23"></line>
+      <line x1="4.22" y1="4.22" x2="5.64" y2="5.64"></line>
+      <line x1="18.36" y1="18.36" x2="19.78" y2="19.78"></line>
+      <line x1="1" y1="12" x2="3" y2="12"></line>
+      <line x1="21" y1="12" x2="23" y2="12"></line>
+      <line x1="4.22" y1="19.78" x2="5.64" y2="18.36"></line>
+      <line x1="18.36" y1="5.64" x2="19.78" y2="4.22"></line>
+    </svg>
+  </symbol>
+  <symbol id="svg-moon" viewBox="0 0 24 24">
+    <title>Dark mode</title>
+    <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
+      stroke-width="1.5" stroke-linecap="round" stroke-linejoin="round" class="icon-tabler-moon">
+      <path stroke="none" d="M0 0h24v24H0z" fill="none" />
+      <path d="M12 3c.132 0 .263 0 .393 0a7.5 7.5 0 0 0 7.92 12.446a9 9 0 1 1 -8.313 -12.454z" />
+    </svg>
+  </symbol>
+  <symbol id="svg-sun-half" viewBox="0 0 24 24">
+    <title>Auto light/dark mode</title>
+    <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
+      stroke-width="1.5" stroke-linecap="round" stroke-linejoin="round" class="icon-tabler-shadow">
+      <path stroke="none" d="M0 0h24v24H0z" fill="none"/>
+      <circle cx="12" cy="12" r="9" />
+      <path d="M13 12h5" />
+      <path d="M13 15h4" />
+      <path d="M13 18h1" />
+      <path d="M13 9h4" />
+      <path d="M13 6h1" />
+    </svg>
+  </symbol>
+</svg>
+
+<input type="checkbox" class="sidebar-toggle" name="__navigation" id="__navigation">
+<input type="checkbox" class="sidebar-toggle" name="__toc" id="__toc">
+<label class="overlay sidebar-overlay" for="__navigation">
+  <div class="visually-hidden">Hide navigation sidebar</div>
+</label>
+<label class="overlay toc-overlay" for="__toc">
+  <div class="visually-hidden">Hide table of contents sidebar</div>
+</label>
+
+
+
+<div class="page">
+  <header class="mobile-header">
+    <div class="header-left">
+      <label class="nav-overlay-icon" for="__navigation">
+        <div class="visually-hidden">Toggle site navigation sidebar</div>
+        <i class="icon"><svg><use href="#svg-menu"></use></svg></i>
+      </label>
+    </div>
+    <div class="header-center">
+      <a href="index.html"><div class="brand">cyto-dl 0.4.0 documentation</div></a>
+    </div>
+    <div class="header-right">
+      <div class="theme-toggle-container theme-toggle-header">
+        <button class="theme-toggle">
+          <div class="visually-hidden">Toggle Light / Dark / Auto color theme</div>
+          <svg class="theme-icon-when-auto"><use href="#svg-sun-half"></use></svg>
+          <svg class="theme-icon-when-dark"><use href="#svg-moon"></use></svg>
+          <svg class="theme-icon-when-light"><use href="#svg-sun"></use></svg>
+        </button>
+      </div>
+      <label class="toc-overlay-icon toc-header-icon" for="__toc">
+        <div class="visually-hidden">Toggle table of contents sidebar</div>
+        <i class="icon"><svg><use href="#svg-toc"></use></svg></i>
+      </label>
+    </div>
+  </header>
+  <aside class="sidebar-drawer">
+    <div class="sidebar-container">
+      
+      <div class="sidebar-sticky"><a class="sidebar-brand" href="index.html">
+  
+  
+  <span class="sidebar-brand-text">cyto-dl 0.4.0 documentation</span>
+  
+</a><form class="sidebar-search-container" method="get" action="search.html" role="search">
+  <input class="sidebar-search" placeholder="Search" name="q" aria-label="Search">
+  <input type="hidden" name="check_keywords" value="yes">
+  <input type="hidden" name="area" value="default">
+</form>
+<div id="searchbox"></div><div class="sidebar-scroll"><div class="sidebar-tree">
+  <p class="caption" role="heading"><span class="caption-text">Contents:</span></p>
+<ul class="current">
+<li class="toctree-l1"><a class="reference internal" href="index.html">Home</a></li>
+<li class="toctree-l1"><a class="reference internal" href="installation.html">Installation</a></li>
+<li class="toctree-l1"><a class="reference internal" href="using_examples.html">Modifying the Example Configs</a></li>
+<li class="toctree-l1"><a class="reference internal" href="api.html">Using the API</a></li>
+<li class="toctree-l1 current has-children"><a class="reference internal" href="modules.html">Package modules</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-1" name="toctree-checkbox-1" role="switch" type="checkbox"/><label for="toctree-checkbox-1"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
+<li class="toctree-l2 current has-children"><a class="reference internal" href="cyto_dl.html">cyto_dl package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-2" name="toctree-checkbox-2" role="switch" type="checkbox"/><label for="toctree-checkbox-2"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.api.html">cyto_dl.api package</a><input class="toctree-checkbox" id="toctree-checkbox-3" name="toctree-checkbox-3" role="switch" type="checkbox"/><label for="toctree-checkbox-3"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.api.cyto_dl_model.html">cyto_dl.api.cyto_dl_model package</a><input class="toctree-checkbox" id="toctree-checkbox-4" name="toctree-checkbox-4" role="switch" type="checkbox"/><label for="toctree-checkbox-4"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.api.cyto_dl_model.cyto_dl_base_model.html">cyto_dl.api.cyto_dl_model.cyto_dl_base_model module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.api.cyto_dl_model.segmentation_plugin_model.html">cyto_dl.api.cyto_dl_model.segmentation_plugin_model module</a></li>
+</ul>
+</li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.api.data.html">cyto_dl.api.data module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.api.model.html">cyto_dl.api.model module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.callbacks.html">cyto_dl.callbacks package</a><input class="toctree-checkbox" id="toctree-checkbox-5" name="toctree-checkbox-5" role="switch" type="checkbox"/><label for="toctree-checkbox-5"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.callbacks.callback_utils.html">cyto_dl.callbacks.callback_utils module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.callbacks.csv_saver.html">cyto_dl.callbacks.csv_saver module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.callbacks.image_saver.html">cyto_dl.callbacks.image_saver module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.callbacks.latent_walk.html">cyto_dl.callbacks.latent_walk module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.callbacks.layer_freeze.html">cyto_dl.callbacks.layer_freeze module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.callbacks.model_utils.html">cyto_dl.callbacks.model_utils module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.callbacks.outlier_detection.html">cyto_dl.callbacks.outlier_detection module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.dataframe.html">cyto_dl.dataframe package</a><input class="toctree-checkbox" id="toctree-checkbox-6" name="toctree-checkbox-6" role="switch" type="checkbox"/><label for="toctree-checkbox-6"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.dataframe.transforms.html">cyto_dl.dataframe.transforms package</a><input class="toctree-checkbox" id="toctree-checkbox-7" name="toctree-checkbox-7" role="switch" type="checkbox"/><label for="toctree-checkbox-7"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.dataframe.transforms.filter.html">cyto_dl.dataframe.transforms.filter module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.dataframe.transforms.group_cols.html">cyto_dl.dataframe.transforms.group_cols module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.dataframe.transforms.misc.html">cyto_dl.dataframe.transforms.misc module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.dataframe.transforms.split.html">cyto_dl.dataframe.transforms.split module</a></li>
+</ul>
+</li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.dataframe.readers.html">cyto_dl.dataframe.readers module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.datamodules.html">cyto_dl.datamodules package</a><input class="toctree-checkbox" id="toctree-checkbox-8" name="toctree-checkbox-8" role="switch" type="checkbox"/><label for="toctree-checkbox-8"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.datamodules.dataframe.html">cyto_dl.datamodules.dataframe package</a><input class="toctree-checkbox" id="toctree-checkbox-9" name="toctree-checkbox-9" role="switch" type="checkbox"/><label for="toctree-checkbox-9"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.datamodules.dataframe.dataframe_datamodule.html">cyto_dl.datamodules.dataframe.dataframe_datamodule module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.datamodules.dataframe.grouped_dataframe_datamodule.html">cyto_dl.datamodules.dataframe.grouped_dataframe_datamodule module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.datamodules.dataframe.utils.html">cyto_dl.datamodules.dataframe.utils module</a></li>
+</ul>
+</li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.datamodules.array.html">cyto_dl.datamodules.array module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.datamodules.data_dict.html">cyto_dl.datamodules.data_dict module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.datamodules.folder.html">cyto_dl.datamodules.folder module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.datamodules.multidim_image.html">cyto_dl.datamodules.multidim_image module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.datamodules.smartcache.html">cyto_dl.datamodules.smartcache module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.datamodules.torchvision.html">cyto_dl.datamodules.torchvision module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.image.html">cyto_dl.image package</a><input class="toctree-checkbox" id="toctree-checkbox-10" name="toctree-checkbox-10" role="switch" type="checkbox"/><label for="toctree-checkbox-10"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.image.io.html">cyto_dl.image.io package</a><input class="toctree-checkbox" id="toctree-checkbox-11" name="toctree-checkbox-11" role="switch" type="checkbox"/><label for="toctree-checkbox-11"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.io.aicsimage_loader.html">cyto_dl.image.io.aicsimage_loader module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.io.monai_bio_reader.html">cyto_dl.image.io.monai_bio_reader module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.io.numpy_reader.html">cyto_dl.image.io.numpy_reader module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.io.ome_zarr_reader.html">cyto_dl.image.io.ome_zarr_reader module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.io.polygon_loader.html">cyto_dl.image.io.polygon_loader module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.io.skimage_reader.html">cyto_dl.image.io.skimage_reader module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.image.transforms.html">cyto_dl.image.transforms package</a><input class="toctree-checkbox" id="toctree-checkbox-12" name="toctree-checkbox-12" role="switch" type="checkbox"/><label for="toctree-checkbox-12"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.add_meta.html">cyto_dl.image.transforms.add_meta module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.bright_sampler.html">cyto_dl.image.transforms.bright_sampler module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.clip.html">cyto_dl.image.transforms.clip module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.contrastadjust.html">cyto_dl.image.transforms.contrastadjust module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.generate_jepa_masks.html">cyto_dl.image.transforms.generate_jepa_masks module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.merge.html">cyto_dl.image.transforms.merge module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.multiscale_cropper.html">cyto_dl.image.transforms.multiscale_cropper module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.pad.html">cyto_dl.image.transforms.pad module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.project.html">cyto_dl.image.transforms.project module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.rotation_mask_transform.html">cyto_dl.image.transforms.rotation_mask_transform module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.save.html">cyto_dl.image.transforms.save module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.so2_random_rotation.html">cyto_dl.image.transforms.so2_random_rotation module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.track_transforms.html">cyto_dl.image.transforms.track_transforms module</a></li>
+</ul>
+</li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.loggers.html">cyto_dl.loggers package</a><input class="toctree-checkbox" id="toctree-checkbox-13" name="toctree-checkbox-13" role="switch" type="checkbox"/><label for="toctree-checkbox-13"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.loggers.mlflow.html">cyto_dl.loggers.mlflow module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.models.html">cyto_dl.models package</a><input class="toctree-checkbox" id="toctree-checkbox-14" name="toctree-checkbox-14" role="switch" type="checkbox"/><label for="toctree-checkbox-14"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.models.classification.html">cyto_dl.models.classification package</a><input class="toctree-checkbox" id="toctree-checkbox-15" name="toctree-checkbox-15" role="switch" type="checkbox"/><label for="toctree-checkbox-15"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.classification.classification.html">cyto_dl.models.classification.classification module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.classification.timepoint_classification.html">cyto_dl.models.classification.timepoint_classification module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.models.contrastive.html">cyto_dl.models.contrastive package</a><input class="toctree-checkbox" id="toctree-checkbox-16" name="toctree-checkbox-16" role="switch" type="checkbox"/><label for="toctree-checkbox-16"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.contrastive.contrastive.html">cyto_dl.models.contrastive.contrastive module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.models.handlers.html">cyto_dl.models.handlers package</a><input class="toctree-checkbox" id="toctree-checkbox-17" name="toctree-checkbox-17" role="switch" type="checkbox"/><label for="toctree-checkbox-17"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.handlers.base_handler.html">cyto_dl.models.handlers.base_handler module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.handlers.image_handler.html">cyto_dl.models.handlers.image_handler module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.handlers.load_image_patch.html">cyto_dl.models.handlers.load_image_patch module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.models.im2im.html">cyto_dl.models.im2im package</a><input class="toctree-checkbox" id="toctree-checkbox-18" name="toctree-checkbox-18" role="switch" type="checkbox"/><label for="toctree-checkbox-18"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.models.im2im.utils.html">cyto_dl.models.im2im.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-19" name="toctree-checkbox-19" role="switch" type="checkbox"/><label for="toctree-checkbox-19"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.models.im2im.utils.postprocessing.html">cyto_dl.models.im2im.utils.postprocessing package</a><input class="toctree-checkbox" id="toctree-checkbox-20" name="toctree-checkbox-20" role="switch" type="checkbox"/><label for="toctree-checkbox-20"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.models.im2im.utils.postprocessing.act_thresh_label.html">cyto_dl.models.im2im.utils.postprocessing.act_thresh_label module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.models.im2im.utils.postprocessing.auto_thresh.html">cyto_dl.models.im2im.utils.postprocessing.auto_thresh module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.models.im2im.utils.postprocessing.dict_to_im.html">cyto_dl.models.im2im.utils.postprocessing.dict_to_im module</a></li>
+</ul>
+</li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.models.im2im.utils.instance_seg.html">cyto_dl.models.im2im.utils.instance_seg module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.models.im2im.utils.noise_annealer.html">cyto_dl.models.im2im.utils.noise_annealer module</a></li>
+</ul>
+</li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.im2im.gan.html">cyto_dl.models.im2im.gan module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.im2im.multi_task.html">cyto_dl.models.im2im.multi_task module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.models.jepa.html">cyto_dl.models.jepa package</a><input class="toctree-checkbox" id="toctree-checkbox-21" name="toctree-checkbox-21" role="switch" type="checkbox"/><label for="toctree-checkbox-21"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.jepa.ijepa.html">cyto_dl.models.jepa.ijepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.jepa.iwm.html">cyto_dl.models.jepa.iwm module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.jepa.jepa_base.html">cyto_dl.models.jepa.jepa_base module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.models.utils.html">cyto_dl.models.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-22" name="toctree-checkbox-22" role="switch" type="checkbox"/><label for="toctree-checkbox-22"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.utils.mlflow.html">cyto_dl.models.utils.mlflow module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.utils.utils.html">cyto_dl.models.utils.utils module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.models.vae.html">cyto_dl.models.vae package</a><input class="toctree-checkbox" id="toctree-checkbox-23" name="toctree-checkbox-23" role="switch" type="checkbox"/><label for="toctree-checkbox-23"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.models.vae.o2_spharm_vae.html">cyto_dl.models.vae.o2_spharm_vae package</a><input class="toctree-checkbox" id="toctree-checkbox-24" name="toctree-checkbox-24" role="switch" type="checkbox"/><label for="toctree-checkbox-24"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.models.vae.o2_spharm_vae.o2_spharm_encoder.html">cyto_dl.models.vae.o2_spharm_vae.o2_spharm_encoder module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.models.vae.o2_spharm_vae.o2_spharm_vae.html">cyto_dl.models.vae.o2_spharm_vae.o2_spharm_vae module</a></li>
+</ul>
+</li>
+<li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.models.vae.priors.html">cyto_dl.models.vae.priors package</a><input class="toctree-checkbox" id="toctree-checkbox-25" name="toctree-checkbox-25" role="switch" type="checkbox"/><label for="toctree-checkbox-25"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.models.vae.priors.abstract_prior.html">cyto_dl.models.vae.priors.abstract_prior module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.models.vae.priors.gaussian.html">cyto_dl.models.vae.priors.gaussian module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.models.vae.priors.identity_prior.html">cyto_dl.models.vae.priors.identity_prior module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.models.vae.priors.joint_prior.html">cyto_dl.models.vae.priors.joint_prior module</a></li>
+</ul>
+</li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.vae.base_vae.html">cyto_dl.models.vae.base_vae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.vae.image_canon_vae.html">cyto_dl.models.vae.image_canon_vae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.vae.image_encoder.html">cyto_dl.models.vae.image_encoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.vae.image_vae.html">cyto_dl.models.vae.image_vae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.vae.latent_loss_vae.html">cyto_dl.models.vae.latent_loss_vae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.vae.point_cloud_vae.html">cyto_dl.models.vae.point_cloud_vae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.vae.tabular_vae.html">cyto_dl.models.vae.tabular_vae module</a></li>
+</ul>
+</li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.models.base_model.html">cyto_dl.models.base_model module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.models.basic_model.html">cyto_dl.models.basic_model module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 current has-children"><a class="reference internal" href="cyto_dl.nn.html">cyto_dl.nn package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-26" name="toctree-checkbox-26" role="switch" type="checkbox"/><label for="toctree-checkbox-26"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.discriminators.html">cyto_dl.nn.discriminators package</a><input class="toctree-checkbox" id="toctree-checkbox-27" name="toctree-checkbox-27" role="switch" type="checkbox"/><label for="toctree-checkbox-27"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.discriminators.multi_scale_discriminator.html">cyto_dl.nn.discriminators.multi_scale_discriminator module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.discriminators.n_layer_discriminator.html">cyto_dl.nn.discriminators.n_layer_discriminator module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.head.html">cyto_dl.nn.head package</a><input class="toctree-checkbox" id="toctree-checkbox-28" name="toctree-checkbox-28" role="switch" type="checkbox"/><label for="toctree-checkbox-28"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.head.base_head.html">cyto_dl.nn.head.base_head module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.head.gan_head.html">cyto_dl.nn.head.gan_head module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.head.gan_head_superres.html">cyto_dl.nn.head.gan_head_superres module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.head.mae_head.html">cyto_dl.nn.head.mae_head module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.head.mask_head.html">cyto_dl.nn.head.mask_head module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.head.res_blocks_head.html">cyto_dl.nn.head.res_blocks_head module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.head.vic_reg.html">cyto_dl.nn.head.vic_reg module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.losses.html">cyto_dl.nn.losses package</a><input class="toctree-checkbox" id="toctree-checkbox-29" name="toctree-checkbox-29" role="switch" type="checkbox"/><label for="toctree-checkbox-29"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.adversarial_loss.html">cyto_dl.nn.losses.adversarial_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.chamfer_loss.html">cyto_dl.nn.losses.chamfer_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.continuous_bernoulli.html">cyto_dl.nn.losses.continuous_bernoulli module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.cosine_loss.html">cyto_dl.nn.losses.cosine_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.gan_loss.html">cyto_dl.nn.losses.gan_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.gaussian_nll_loss.html">cyto_dl.nn.losses.gaussian_nll_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.geomloss.html">cyto_dl.nn.losses.geomloss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.loss_wrapper.html">cyto_dl.nn.losses.loss_wrapper module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.spharm_loss.html">cyto_dl.nn.losses.spharm_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.threshold_loss.html">cyto_dl.nn.losses.threshold_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.vic_reg.html">cyto_dl.nn.losses.vic_reg module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.weibull.html">cyto_dl.nn.losses.weibull module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.weighted_mse_loss.html">cyto_dl.nn.losses.weighted_mse_loss module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.point_cloud.html">cyto_dl.nn.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-30" name="toctree-checkbox-30" role="switch" type="checkbox"/><label for="toctree-checkbox-30"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.point_cloud.dgcnn.html">cyto_dl.nn.point_cloud.dgcnn module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.point_cloud.folding_net.html">cyto_dl.nn.point_cloud.folding_net module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.point_cloud.graph_functions.html">cyto_dl.nn.point_cloud.graph_functions module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.point_cloud.vnn.html">cyto_dl.nn.point_cloud.vnn module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 current has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
+<li class="toctree-l5 current has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
+<li class="toctree-l6 current has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7 current current-page"><a class="current reference internal" href="#">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
+</ul>
+</li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
+</ul>
+</li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.hr_skip.html">cyto_dl.nn.hr_skip module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.mlp.html">cyto_dl.nn.mlp module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.res_unit.html">cyto_dl.nn.res_unit module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.spatial_transformer.html">cyto_dl.nn.spatial_transformer module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
+</ul>
+</li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
+</ul>
+</li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.utils.array.html">cyto_dl.utils.array module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.utils.checkpoint.html">cyto_dl.utils.checkpoint module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.utils.config.html">cyto_dl.utils.config module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.utils.download_test_data.html">cyto_dl.utils.download_test_data module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.utils.dummy_dataset.html">cyto_dl.utils.dummy_dataset module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.utils.pylogger.html">cyto_dl.utils.pylogger module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.utils.rich_utils.html">cyto_dl.utils.rich_utils module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.utils.rotation.html">cyto_dl.utils.rotation module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.utils.template_utils.html">cyto_dl.utils.template_utils module</a></li>
+</ul>
+</li>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.compile.html">cyto_dl.compile module</a></li>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.eval.html">cyto_dl.eval module</a></li>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.train.html">cyto_dl.train module</a></li>
+</ul>
+</li>
+</ul>
+</li>
+</ul>
+
+</div>
+</div>
+
+      </div>
+      
+    </div>
+  </aside>
+  <div class="main">
+    <div class="content">
+      <div class="article-container">
+        <a href="#" class="back-to-top muted-link">
+          <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24">
+            <path d="M13 20h-2V8l-5.5 5.5-1.42-1.42L12 4.16l7.92 7.92-1.42 1.42L13 8v12z"></path>
+          </svg>
+          <span>Back to top</span>
+        </a>
+        <div class="content-icon-container">
+          
+<div class="theme-toggle-container theme-toggle-content">
+            <button class="theme-toggle">
+              <div class="visually-hidden">Toggle Light / Dark / Auto color theme</div>
+              <svg class="theme-icon-when-auto"><use href="#svg-sun-half"></use></svg>
+              <svg class="theme-icon-when-dark"><use href="#svg-moon"></use></svg>
+              <svg class="theme-icon-when-light"><use href="#svg-sun"></use></svg>
+            </button>
+          </div>
+          <label class="toc-overlay-icon toc-content-icon" for="__toc">
+            <div class="visually-hidden">Toggle table of contents sidebar</div>
+            <i class="icon"><svg><use href="#svg-toc"></use></svg></i>
+          </label>
+        </div>
+        <article role="main">
+          <section id="module-cyto_dl.nn.vits.blocks.patchify.patchify_hiera">
+<span id="cyto-dl-nn-vits-blocks-patchify-patchify-hiera-module"></span><h1>cyto_dl.nn.vits.blocks.patchify.patchify_hiera module<a class="headerlink" href="#module-cyto_dl.nn.vits.blocks.patchify.patchify_hiera" title="Permalink to this heading">#</a></h1>
+<dl class="py class">
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.PatchifyHiera">
+<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">cyto_dl.nn.vits.blocks.patchify.patchify_hiera.</span></span><span class="sig-name descname"><span class="pre">PatchifyHiera</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">patch_size</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">]</span></span></span></em>, <em class="sig-param"><span class="n"><span class="pre">n_patches</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">]</span></span></span></em>, <em class="sig-param"><span class="n"><span class="pre">emb_dim</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">64</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">spatial_dims</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">3</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">context_pixels</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">]</span></span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">[0,</span> <span class="pre">0,</span> <span class="pre">0]</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">input_channels</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">1</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">tasks</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">str</span><span class="p"><span class="pre">]</span></span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">[]</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">mask_units_per_dim</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">]</span></span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">[8,</span> <span class="pre">8,</span> <span class="pre">8]</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/blocks/patchify/patchify_hiera.html#PatchifyHiera"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.blocks.patchify.patchify_hiera.PatchifyHiera" title="Permalink to this definition">#</a></dt>
+<dd><p>Bases: <a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html#cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase" title="cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase"><code class="xref py py-class docutils literal notranslate"><span class="pre">PatchifyBase</span></code></a></p>
+<p>Class for converting images to a sequence of patches with positional embeddings, masked at
+the level of mask units (groups of patches specified by mask_units_per_dim).</p>
+<dl class="simple">
+<dt>patch_size: List[int]</dt><dd><p>Size of each patch in pix (ZYX order for 3D, YX order for 2D)</p>
+</dd>
+<dt>n_patches: List[int]</dt><dd><p>Number of patches in each spatial dimension (ZYX order for 3D, YX order for 2D)</p>
+</dd>
+<dt>emb_dim: int</dt><dd><p>Dimension of encoder</p>
+</dd>
+<dt>spatial_dims: int</dt><dd><p>Number of spatial dimensions</p>
+</dd>
+<dt>context_pixels: List[int]</dt><dd><p>Number of extra pixels around each patch to include in convolutional embedding to encoder dimension.</p>
+</dd>
+<dt>input_channels: int</dt><dd><p>Number of input channels</p>
+</dd>
+<dt>tasks: List[str]</dt><dd><p>List of tasks to encode</p>
+</dd>
+<dt>mask_units_per_dim: List[int]</dt><dd><p>Number of mask units in each spatial dimension (ZYX order for 3D, YX order for 2D)</p>
+</dd>
+</dl>
+<dl class="py method">
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.PatchifyHiera.create_img2token">
+<span class="sig-name descname"><span class="pre">create_img2token</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">mask_units_per_dim</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/blocks/patchify/patchify_hiera.html#PatchifyHiera.create_img2token"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.blocks.patchify.patchify_hiera.PatchifyHiera.create_img2token" title="Permalink to this definition">#</a></dt>
+<dd></dd></dl>
+
+<dl class="py method">
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.PatchifyHiera.extract_visible_tokens">
+<span class="sig-name descname"><span class="pre">extract_visible_tokens</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">tokens</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">forward_indexes</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">n_visible_patches</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/blocks/patchify/patchify_hiera.html#PatchifyHiera.extract_visible_tokens"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.blocks.patchify.patchify_hiera.PatchifyHiera.extract_visible_tokens" title="Permalink to this definition">#</a></dt>
+<dd></dd></dl>
+
+<dl class="py method">
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.PatchifyHiera.get_mask_args">
+<span class="sig-name descname"><span class="pre">get_mask_args</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">mask_ratio</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/blocks/patchify/patchify_hiera.html#PatchifyHiera.get_mask_args"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.blocks.patchify.patchify_hiera.PatchifyHiera.get_mask_args" title="Permalink to this definition">#</a></dt>
+<dd></dd></dl>
+
+<dl class="py property">
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.PatchifyHiera.img2token">
+<em class="property"><span class="pre">property</span><span class="w"> </span></em><span class="sig-name descname"><span class="pre">img2token</span></span><a class="headerlink" href="#cyto_dl.nn.vits.blocks.patchify.patchify_hiera.PatchifyHiera.img2token" title="Permalink to this definition">#</a></dt>
+<dd></dd></dl>
+
+</dd></dl>
+
+<dl class="py function">
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.take_indexes_mask">
+<span class="sig-prename descclassname"><span class="pre">cyto_dl.nn.vits.blocks.patchify.patchify_hiera.</span></span><span class="sig-name descname"><span class="pre">take_indexes_mask</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">sequences</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">indexes</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/blocks/patchify/patchify_hiera.html#take_indexes_mask"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.blocks.patchify.patchify_hiera.take_indexes_mask" title="Permalink to this definition">#</a></dt>
+<dd><p>sequences: batch x mask units x patches x emb_dim
+indexes: mask_units x batch</p>
+</dd></dl>
+
+</section>
+
+        </article>
+      </div>
+      <footer>
+        
+        <div class="related-pages">
+          <a class="next-page" href="cyto_dl.nn.vits.blocks.cross_attention.html">
+              <div class="page-info">
+                <div class="context">
+                  <span>Next</span>
+                </div>
+                <div class="title">cyto_dl.nn.vits.blocks.cross_attention module</div>
+              </div>
+              <svg class="furo-related-icon"><use href="#svg-arrow-right"></use></svg>
+            </a>
+          <a class="prev-page" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">
+              <svg class="furo-related-icon"><use href="#svg-arrow-right"></use></svg>
+              <div class="page-info">
+                <div class="context">
+                  <span>Previous</span>
+                </div>
+                
+                <div class="title">cyto_dl.nn.vits.blocks.patchify.patchify_base module</div>
+                
+              </div>
+            </a>
+        </div>
+        <div class="bottom-of-page">
+          <div class="left-details">
+            <div class="copyright">
+                Copyright &#169; 2023
+            </div>
+            Made with <a href="https://www.sphinx-doc.org/">Sphinx</a> and <a class="muted-link" href="https://pradyunsg.me">@pradyunsg</a>'s
+            
+            <a href="https://github.com/pradyunsg/furo">Furo</a>
+            
+          </div>
+          <div class="right-details">
+            
+          </div>
+        </div>
+        
+      </footer>
+    </div>
+    <aside class="toc-drawer">
+      
+      
+      <div class="toc-sticky toc-scroll">
+        <div class="toc-title-container">
+          <span class="toc-title">
+            On this page
+          </span>
+        </div>
+        <div class="toc-tree-container">
+          <div class="toc-tree">
+            <ul>
+<li><a class="reference internal" href="#">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a><ul>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.blocks.patchify.patchify_hiera.PatchifyHiera"><code class="docutils literal notranslate"><span class="pre">PatchifyHiera</span></code></a><ul>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.blocks.patchify.patchify_hiera.PatchifyHiera.create_img2token"><code class="docutils literal notranslate"><span class="pre">PatchifyHiera.create_img2token()</span></code></a></li>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.blocks.patchify.patchify_hiera.PatchifyHiera.extract_visible_tokens"><code class="docutils literal notranslate"><span class="pre">PatchifyHiera.extract_visible_tokens()</span></code></a></li>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.blocks.patchify.patchify_hiera.PatchifyHiera.get_mask_args"><code class="docutils literal notranslate"><span class="pre">PatchifyHiera.get_mask_args()</span></code></a></li>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.blocks.patchify.patchify_hiera.PatchifyHiera.img2token"><code class="docutils literal notranslate"><span class="pre">PatchifyHiera.img2token</span></code></a></li>
+</ul>
+</li>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.blocks.patchify.patchify_hiera.take_indexes_mask"><code class="docutils literal notranslate"><span class="pre">take_indexes_mask()</span></code></a></li>
+</ul>
+</li>
+</ul>
+
+          </div>
+        </div>
+      </div>
+      
+      
+    </aside>
+  </div>
+</div><script data-url_root="./" id="documentation_options" src="_static/documentation_options.js"></script>
+    <script src="_static/doctools.js"></script>
+    <script src="_static/sphinx_highlight.js"></script>
+    <script src="_static/scripts/furo.js"></script>
+    </body>
+</html>
\ No newline at end of file
diff --git a/cyto_dl.nn.vits.decoder.html b/cyto_dl.nn.vits.decoder.html
new file mode 100644
index 000000000..20fe7f43b
--- /dev/null
+++ b/cyto_dl.nn.vits.decoder.html
@@ -0,0 +1,598 @@
+<!doctype html>
+<html class="no-js" lang="en">
+  <head><meta charset="utf-8"/>
+    <meta name="viewport" content="width=device-width,initial-scale=1"/>
+    <meta name="color-scheme" content="light dark"><meta name="generator" content="Docutils 0.19: https://docutils.sourceforge.io/" />
+<link rel="index" title="Index" href="genindex.html" /><link rel="search" title="Search" href="search.html" /><link rel="next" title="cyto_dl.nn.vits.encoder module" href="cyto_dl.nn.vits.encoder.html" /><link rel="prev" title="cyto_dl.nn.vits.blocks.masked_unit_attention module" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html" />
+
+    <!-- Generated with Sphinx 6.1.3 and Furo 2023.03.23 -->
+        <title>cyto_dl.nn.vits.decoder module - cyto-dl 0.4.0 documentation</title>
+      <link rel="stylesheet" type="text/css" href="_static/pygments.css" />
+    <link rel="stylesheet" type="text/css" href="_static/styles/furo.css?digest=fad236701ea90a88636c2a8c73b44ae642ed2a53" />
+    <link rel="stylesheet" type="text/css" href="_static/styles/furo-extensions.css?digest=30d1aed668e5c3a91c3e3bf6a60b675221979f0e" />
+    
+    
+
+
+<style>
+  body {
+    --color-code-background: #eeffcc;
+  --color-code-foreground: black;
+  
+  }
+  @media not print {
+    body[data-theme="dark"] {
+      --color-code-background: #202020;
+  --color-code-foreground: #d0d0d0;
+  
+    }
+    @media (prefers-color-scheme: dark) {
+      body:not([data-theme="light"]) {
+        --color-code-background: #202020;
+  --color-code-foreground: #d0d0d0;
+  
+      }
+    }
+  }
+</style></head>
+  <body>
+    
+    <script>
+      document.body.dataset.theme = localStorage.getItem("theme") || "auto";
+    </script>
+    
+
+<svg xmlns="http://www.w3.org/2000/svg" style="display: none;">
+  <symbol id="svg-toc" viewBox="0 0 24 24">
+    <title>Contents</title>
+    <svg stroke="currentColor" fill="currentColor" stroke-width="0" viewBox="0 0 1024 1024">
+      <path d="M408 442h480c4.4 0 8-3.6 8-8v-56c0-4.4-3.6-8-8-8H408c-4.4 0-8 3.6-8 8v56c0 4.4 3.6 8 8 8zm-8 204c0 4.4 3.6 8 8 8h480c4.4 0 8-3.6 8-8v-56c0-4.4-3.6-8-8-8H408c-4.4 0-8 3.6-8 8v56zm504-486H120c-4.4 0-8 3.6-8 8v56c0 4.4 3.6 8 8 8h784c4.4 0 8-3.6 8-8v-56c0-4.4-3.6-8-8-8zm0 632H120c-4.4 0-8 3.6-8 8v56c0 4.4 3.6 8 8 8h784c4.4 0 8-3.6 8-8v-56c0-4.4-3.6-8-8-8zM115.4 518.9L271.7 642c5.8 4.6 14.4.5 14.4-6.9V388.9c0-7.4-8.5-11.5-14.4-6.9L115.4 505.1a8.74 8.74 0 0 0 0 13.8z"/>
+    </svg>
+  </symbol>
+  <symbol id="svg-menu" viewBox="0 0 24 24">
+    <title>Menu</title>
+    <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
+      stroke-width="2" stroke-linecap="round" stroke-linejoin="round" class="feather-menu">
+      <line x1="3" y1="12" x2="21" y2="12"></line>
+      <line x1="3" y1="6" x2="21" y2="6"></line>
+      <line x1="3" y1="18" x2="21" y2="18"></line>
+    </svg>
+  </symbol>
+  <symbol id="svg-arrow-right" viewBox="0 0 24 24">
+    <title>Expand</title>
+    <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
+      stroke-width="2" stroke-linecap="round" stroke-linejoin="round" class="feather-chevron-right">
+      <polyline points="9 18 15 12 9 6"></polyline>
+    </svg>
+  </symbol>
+  <symbol id="svg-sun" viewBox="0 0 24 24">
+    <title>Light mode</title>
+    <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
+      stroke-width="1.5" stroke-linecap="round" stroke-linejoin="round" class="feather-sun">
+      <circle cx="12" cy="12" r="5"></circle>
+      <line x1="12" y1="1" x2="12" y2="3"></line>
+      <line x1="12" y1="21" x2="12" y2="23"></line>
+      <line x1="4.22" y1="4.22" x2="5.64" y2="5.64"></line>
+      <line x1="18.36" y1="18.36" x2="19.78" y2="19.78"></line>
+      <line x1="1" y1="12" x2="3" y2="12"></line>
+      <line x1="21" y1="12" x2="23" y2="12"></line>
+      <line x1="4.22" y1="19.78" x2="5.64" y2="18.36"></line>
+      <line x1="18.36" y1="5.64" x2="19.78" y2="4.22"></line>
+    </svg>
+  </symbol>
+  <symbol id="svg-moon" viewBox="0 0 24 24">
+    <title>Dark mode</title>
+    <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
+      stroke-width="1.5" stroke-linecap="round" stroke-linejoin="round" class="icon-tabler-moon">
+      <path stroke="none" d="M0 0h24v24H0z" fill="none" />
+      <path d="M12 3c.132 0 .263 0 .393 0a7.5 7.5 0 0 0 7.92 12.446a9 9 0 1 1 -8.313 -12.454z" />
+    </svg>
+  </symbol>
+  <symbol id="svg-sun-half" viewBox="0 0 24 24">
+    <title>Auto light/dark mode</title>
+    <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
+      stroke-width="1.5" stroke-linecap="round" stroke-linejoin="round" class="icon-tabler-shadow">
+      <path stroke="none" d="M0 0h24v24H0z" fill="none"/>
+      <circle cx="12" cy="12" r="9" />
+      <path d="M13 12h5" />
+      <path d="M13 15h4" />
+      <path d="M13 18h1" />
+      <path d="M13 9h4" />
+      <path d="M13 6h1" />
+    </svg>
+  </symbol>
+</svg>
+
+<input type="checkbox" class="sidebar-toggle" name="__navigation" id="__navigation">
+<input type="checkbox" class="sidebar-toggle" name="__toc" id="__toc">
+<label class="overlay sidebar-overlay" for="__navigation">
+  <div class="visually-hidden">Hide navigation sidebar</div>
+</label>
+<label class="overlay toc-overlay" for="__toc">
+  <div class="visually-hidden">Hide table of contents sidebar</div>
+</label>
+
+
+
+<div class="page">
+  <header class="mobile-header">
+    <div class="header-left">
+      <label class="nav-overlay-icon" for="__navigation">
+        <div class="visually-hidden">Toggle site navigation sidebar</div>
+        <i class="icon"><svg><use href="#svg-menu"></use></svg></i>
+      </label>
+    </div>
+    <div class="header-center">
+      <a href="index.html"><div class="brand">cyto-dl 0.4.0 documentation</div></a>
+    </div>
+    <div class="header-right">
+      <div class="theme-toggle-container theme-toggle-header">
+        <button class="theme-toggle">
+          <div class="visually-hidden">Toggle Light / Dark / Auto color theme</div>
+          <svg class="theme-icon-when-auto"><use href="#svg-sun-half"></use></svg>
+          <svg class="theme-icon-when-dark"><use href="#svg-moon"></use></svg>
+          <svg class="theme-icon-when-light"><use href="#svg-sun"></use></svg>
+        </button>
+      </div>
+      <label class="toc-overlay-icon toc-header-icon" for="__toc">
+        <div class="visually-hidden">Toggle table of contents sidebar</div>
+        <i class="icon"><svg><use href="#svg-toc"></use></svg></i>
+      </label>
+    </div>
+  </header>
+  <aside class="sidebar-drawer">
+    <div class="sidebar-container">
+      
+      <div class="sidebar-sticky"><a class="sidebar-brand" href="index.html">
+  
+  
+  <span class="sidebar-brand-text">cyto-dl 0.4.0 documentation</span>
+  
+</a><form class="sidebar-search-container" method="get" action="search.html" role="search">
+  <input class="sidebar-search" placeholder="Search" name="q" aria-label="Search">
+  <input type="hidden" name="check_keywords" value="yes">
+  <input type="hidden" name="area" value="default">
+</form>
+<div id="searchbox"></div><div class="sidebar-scroll"><div class="sidebar-tree">
+  <p class="caption" role="heading"><span class="caption-text">Contents:</span></p>
+<ul class="current">
+<li class="toctree-l1"><a class="reference internal" href="index.html">Home</a></li>
+<li class="toctree-l1"><a class="reference internal" href="installation.html">Installation</a></li>
+<li class="toctree-l1"><a class="reference internal" href="using_examples.html">Modifying the Example Configs</a></li>
+<li class="toctree-l1"><a class="reference internal" href="api.html">Using the API</a></li>
+<li class="toctree-l1 current has-children"><a class="reference internal" href="modules.html">Package modules</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-1" name="toctree-checkbox-1" role="switch" type="checkbox"/><label for="toctree-checkbox-1"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
+<li class="toctree-l2 current has-children"><a class="reference internal" href="cyto_dl.html">cyto_dl package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-2" name="toctree-checkbox-2" role="switch" type="checkbox"/><label for="toctree-checkbox-2"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.api.html">cyto_dl.api package</a><input class="toctree-checkbox" id="toctree-checkbox-3" name="toctree-checkbox-3" role="switch" type="checkbox"/><label for="toctree-checkbox-3"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.api.cyto_dl_model.html">cyto_dl.api.cyto_dl_model package</a><input class="toctree-checkbox" id="toctree-checkbox-4" name="toctree-checkbox-4" role="switch" type="checkbox"/><label for="toctree-checkbox-4"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.api.cyto_dl_model.cyto_dl_base_model.html">cyto_dl.api.cyto_dl_model.cyto_dl_base_model module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.api.cyto_dl_model.segmentation_plugin_model.html">cyto_dl.api.cyto_dl_model.segmentation_plugin_model module</a></li>
+</ul>
+</li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.api.data.html">cyto_dl.api.data module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.api.model.html">cyto_dl.api.model module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.callbacks.html">cyto_dl.callbacks package</a><input class="toctree-checkbox" id="toctree-checkbox-5" name="toctree-checkbox-5" role="switch" type="checkbox"/><label for="toctree-checkbox-5"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.callbacks.callback_utils.html">cyto_dl.callbacks.callback_utils module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.callbacks.csv_saver.html">cyto_dl.callbacks.csv_saver module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.callbacks.image_saver.html">cyto_dl.callbacks.image_saver module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.callbacks.latent_walk.html">cyto_dl.callbacks.latent_walk module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.callbacks.layer_freeze.html">cyto_dl.callbacks.layer_freeze module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.callbacks.model_utils.html">cyto_dl.callbacks.model_utils module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.callbacks.outlier_detection.html">cyto_dl.callbacks.outlier_detection module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.dataframe.html">cyto_dl.dataframe package</a><input class="toctree-checkbox" id="toctree-checkbox-6" name="toctree-checkbox-6" role="switch" type="checkbox"/><label for="toctree-checkbox-6"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.dataframe.transforms.html">cyto_dl.dataframe.transforms package</a><input class="toctree-checkbox" id="toctree-checkbox-7" name="toctree-checkbox-7" role="switch" type="checkbox"/><label for="toctree-checkbox-7"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.dataframe.transforms.filter.html">cyto_dl.dataframe.transforms.filter module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.dataframe.transforms.group_cols.html">cyto_dl.dataframe.transforms.group_cols module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.dataframe.transforms.misc.html">cyto_dl.dataframe.transforms.misc module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.dataframe.transforms.split.html">cyto_dl.dataframe.transforms.split module</a></li>
+</ul>
+</li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.dataframe.readers.html">cyto_dl.dataframe.readers module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.datamodules.html">cyto_dl.datamodules package</a><input class="toctree-checkbox" id="toctree-checkbox-8" name="toctree-checkbox-8" role="switch" type="checkbox"/><label for="toctree-checkbox-8"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.datamodules.dataframe.html">cyto_dl.datamodules.dataframe package</a><input class="toctree-checkbox" id="toctree-checkbox-9" name="toctree-checkbox-9" role="switch" type="checkbox"/><label for="toctree-checkbox-9"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.datamodules.dataframe.dataframe_datamodule.html">cyto_dl.datamodules.dataframe.dataframe_datamodule module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.datamodules.dataframe.grouped_dataframe_datamodule.html">cyto_dl.datamodules.dataframe.grouped_dataframe_datamodule module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.datamodules.dataframe.utils.html">cyto_dl.datamodules.dataframe.utils module</a></li>
+</ul>
+</li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.datamodules.array.html">cyto_dl.datamodules.array module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.datamodules.data_dict.html">cyto_dl.datamodules.data_dict module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.datamodules.folder.html">cyto_dl.datamodules.folder module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.datamodules.multidim_image.html">cyto_dl.datamodules.multidim_image module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.datamodules.smartcache.html">cyto_dl.datamodules.smartcache module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.datamodules.torchvision.html">cyto_dl.datamodules.torchvision module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.image.html">cyto_dl.image package</a><input class="toctree-checkbox" id="toctree-checkbox-10" name="toctree-checkbox-10" role="switch" type="checkbox"/><label for="toctree-checkbox-10"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.image.io.html">cyto_dl.image.io package</a><input class="toctree-checkbox" id="toctree-checkbox-11" name="toctree-checkbox-11" role="switch" type="checkbox"/><label for="toctree-checkbox-11"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.io.aicsimage_loader.html">cyto_dl.image.io.aicsimage_loader module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.io.monai_bio_reader.html">cyto_dl.image.io.monai_bio_reader module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.io.numpy_reader.html">cyto_dl.image.io.numpy_reader module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.io.ome_zarr_reader.html">cyto_dl.image.io.ome_zarr_reader module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.io.polygon_loader.html">cyto_dl.image.io.polygon_loader module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.io.skimage_reader.html">cyto_dl.image.io.skimage_reader module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.image.transforms.html">cyto_dl.image.transforms package</a><input class="toctree-checkbox" id="toctree-checkbox-12" name="toctree-checkbox-12" role="switch" type="checkbox"/><label for="toctree-checkbox-12"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.add_meta.html">cyto_dl.image.transforms.add_meta module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.bright_sampler.html">cyto_dl.image.transforms.bright_sampler module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.clip.html">cyto_dl.image.transforms.clip module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.contrastadjust.html">cyto_dl.image.transforms.contrastadjust module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.generate_jepa_masks.html">cyto_dl.image.transforms.generate_jepa_masks module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.merge.html">cyto_dl.image.transforms.merge module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.multiscale_cropper.html">cyto_dl.image.transforms.multiscale_cropper module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.pad.html">cyto_dl.image.transforms.pad module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.project.html">cyto_dl.image.transforms.project module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.rotation_mask_transform.html">cyto_dl.image.transforms.rotation_mask_transform module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.save.html">cyto_dl.image.transforms.save module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.so2_random_rotation.html">cyto_dl.image.transforms.so2_random_rotation module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.track_transforms.html">cyto_dl.image.transforms.track_transforms module</a></li>
+</ul>
+</li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.loggers.html">cyto_dl.loggers package</a><input class="toctree-checkbox" id="toctree-checkbox-13" name="toctree-checkbox-13" role="switch" type="checkbox"/><label for="toctree-checkbox-13"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.loggers.mlflow.html">cyto_dl.loggers.mlflow module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.models.html">cyto_dl.models package</a><input class="toctree-checkbox" id="toctree-checkbox-14" name="toctree-checkbox-14" role="switch" type="checkbox"/><label for="toctree-checkbox-14"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.models.classification.html">cyto_dl.models.classification package</a><input class="toctree-checkbox" id="toctree-checkbox-15" name="toctree-checkbox-15" role="switch" type="checkbox"/><label for="toctree-checkbox-15"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.classification.classification.html">cyto_dl.models.classification.classification module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.classification.timepoint_classification.html">cyto_dl.models.classification.timepoint_classification module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.models.contrastive.html">cyto_dl.models.contrastive package</a><input class="toctree-checkbox" id="toctree-checkbox-16" name="toctree-checkbox-16" role="switch" type="checkbox"/><label for="toctree-checkbox-16"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.contrastive.contrastive.html">cyto_dl.models.contrastive.contrastive module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.models.handlers.html">cyto_dl.models.handlers package</a><input class="toctree-checkbox" id="toctree-checkbox-17" name="toctree-checkbox-17" role="switch" type="checkbox"/><label for="toctree-checkbox-17"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.handlers.base_handler.html">cyto_dl.models.handlers.base_handler module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.handlers.image_handler.html">cyto_dl.models.handlers.image_handler module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.handlers.load_image_patch.html">cyto_dl.models.handlers.load_image_patch module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.models.im2im.html">cyto_dl.models.im2im package</a><input class="toctree-checkbox" id="toctree-checkbox-18" name="toctree-checkbox-18" role="switch" type="checkbox"/><label for="toctree-checkbox-18"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.models.im2im.utils.html">cyto_dl.models.im2im.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-19" name="toctree-checkbox-19" role="switch" type="checkbox"/><label for="toctree-checkbox-19"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.models.im2im.utils.postprocessing.html">cyto_dl.models.im2im.utils.postprocessing package</a><input class="toctree-checkbox" id="toctree-checkbox-20" name="toctree-checkbox-20" role="switch" type="checkbox"/><label for="toctree-checkbox-20"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.models.im2im.utils.postprocessing.act_thresh_label.html">cyto_dl.models.im2im.utils.postprocessing.act_thresh_label module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.models.im2im.utils.postprocessing.auto_thresh.html">cyto_dl.models.im2im.utils.postprocessing.auto_thresh module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.models.im2im.utils.postprocessing.dict_to_im.html">cyto_dl.models.im2im.utils.postprocessing.dict_to_im module</a></li>
+</ul>
+</li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.models.im2im.utils.instance_seg.html">cyto_dl.models.im2im.utils.instance_seg module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.models.im2im.utils.noise_annealer.html">cyto_dl.models.im2im.utils.noise_annealer module</a></li>
+</ul>
+</li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.im2im.gan.html">cyto_dl.models.im2im.gan module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.im2im.multi_task.html">cyto_dl.models.im2im.multi_task module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.models.jepa.html">cyto_dl.models.jepa package</a><input class="toctree-checkbox" id="toctree-checkbox-21" name="toctree-checkbox-21" role="switch" type="checkbox"/><label for="toctree-checkbox-21"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.jepa.ijepa.html">cyto_dl.models.jepa.ijepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.jepa.iwm.html">cyto_dl.models.jepa.iwm module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.jepa.jepa_base.html">cyto_dl.models.jepa.jepa_base module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.models.utils.html">cyto_dl.models.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-22" name="toctree-checkbox-22" role="switch" type="checkbox"/><label for="toctree-checkbox-22"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.utils.mlflow.html">cyto_dl.models.utils.mlflow module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.utils.utils.html">cyto_dl.models.utils.utils module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.models.vae.html">cyto_dl.models.vae package</a><input class="toctree-checkbox" id="toctree-checkbox-23" name="toctree-checkbox-23" role="switch" type="checkbox"/><label for="toctree-checkbox-23"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.models.vae.o2_spharm_vae.html">cyto_dl.models.vae.o2_spharm_vae package</a><input class="toctree-checkbox" id="toctree-checkbox-24" name="toctree-checkbox-24" role="switch" type="checkbox"/><label for="toctree-checkbox-24"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.models.vae.o2_spharm_vae.o2_spharm_encoder.html">cyto_dl.models.vae.o2_spharm_vae.o2_spharm_encoder module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.models.vae.o2_spharm_vae.o2_spharm_vae.html">cyto_dl.models.vae.o2_spharm_vae.o2_spharm_vae module</a></li>
+</ul>
+</li>
+<li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.models.vae.priors.html">cyto_dl.models.vae.priors package</a><input class="toctree-checkbox" id="toctree-checkbox-25" name="toctree-checkbox-25" role="switch" type="checkbox"/><label for="toctree-checkbox-25"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.models.vae.priors.abstract_prior.html">cyto_dl.models.vae.priors.abstract_prior module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.models.vae.priors.gaussian.html">cyto_dl.models.vae.priors.gaussian module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.models.vae.priors.identity_prior.html">cyto_dl.models.vae.priors.identity_prior module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.models.vae.priors.joint_prior.html">cyto_dl.models.vae.priors.joint_prior module</a></li>
+</ul>
+</li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.vae.base_vae.html">cyto_dl.models.vae.base_vae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.vae.image_canon_vae.html">cyto_dl.models.vae.image_canon_vae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.vae.image_encoder.html">cyto_dl.models.vae.image_encoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.vae.image_vae.html">cyto_dl.models.vae.image_vae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.vae.latent_loss_vae.html">cyto_dl.models.vae.latent_loss_vae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.vae.point_cloud_vae.html">cyto_dl.models.vae.point_cloud_vae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.vae.tabular_vae.html">cyto_dl.models.vae.tabular_vae module</a></li>
+</ul>
+</li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.models.base_model.html">cyto_dl.models.base_model module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.models.basic_model.html">cyto_dl.models.basic_model module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 current has-children"><a class="reference internal" href="cyto_dl.nn.html">cyto_dl.nn package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-26" name="toctree-checkbox-26" role="switch" type="checkbox"/><label for="toctree-checkbox-26"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.discriminators.html">cyto_dl.nn.discriminators package</a><input class="toctree-checkbox" id="toctree-checkbox-27" name="toctree-checkbox-27" role="switch" type="checkbox"/><label for="toctree-checkbox-27"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.discriminators.multi_scale_discriminator.html">cyto_dl.nn.discriminators.multi_scale_discriminator module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.discriminators.n_layer_discriminator.html">cyto_dl.nn.discriminators.n_layer_discriminator module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.head.html">cyto_dl.nn.head package</a><input class="toctree-checkbox" id="toctree-checkbox-28" name="toctree-checkbox-28" role="switch" type="checkbox"/><label for="toctree-checkbox-28"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.head.base_head.html">cyto_dl.nn.head.base_head module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.head.gan_head.html">cyto_dl.nn.head.gan_head module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.head.gan_head_superres.html">cyto_dl.nn.head.gan_head_superres module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.head.mae_head.html">cyto_dl.nn.head.mae_head module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.head.mask_head.html">cyto_dl.nn.head.mask_head module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.head.res_blocks_head.html">cyto_dl.nn.head.res_blocks_head module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.head.vic_reg.html">cyto_dl.nn.head.vic_reg module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.losses.html">cyto_dl.nn.losses package</a><input class="toctree-checkbox" id="toctree-checkbox-29" name="toctree-checkbox-29" role="switch" type="checkbox"/><label for="toctree-checkbox-29"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.adversarial_loss.html">cyto_dl.nn.losses.adversarial_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.chamfer_loss.html">cyto_dl.nn.losses.chamfer_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.continuous_bernoulli.html">cyto_dl.nn.losses.continuous_bernoulli module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.cosine_loss.html">cyto_dl.nn.losses.cosine_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.gan_loss.html">cyto_dl.nn.losses.gan_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.gaussian_nll_loss.html">cyto_dl.nn.losses.gaussian_nll_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.geomloss.html">cyto_dl.nn.losses.geomloss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.loss_wrapper.html">cyto_dl.nn.losses.loss_wrapper module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.spharm_loss.html">cyto_dl.nn.losses.spharm_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.threshold_loss.html">cyto_dl.nn.losses.threshold_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.vic_reg.html">cyto_dl.nn.losses.vic_reg module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.weibull.html">cyto_dl.nn.losses.weibull module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.weighted_mse_loss.html">cyto_dl.nn.losses.weighted_mse_loss module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.point_cloud.html">cyto_dl.nn.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-30" name="toctree-checkbox-30" role="switch" type="checkbox"/><label for="toctree-checkbox-30"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.point_cloud.dgcnn.html">cyto_dl.nn.point_cloud.dgcnn module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.point_cloud.folding_net.html">cyto_dl.nn.point_cloud.folding_net module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.point_cloud.graph_functions.html">cyto_dl.nn.point_cloud.graph_functions module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.point_cloud.vnn.html">cyto_dl.nn.point_cloud.vnn module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 current has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
+<li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
+</ul>
+</li>
+<li class="toctree-l5 current current-page"><a class="current reference internal" href="#">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
+</ul>
+</li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.hr_skip.html">cyto_dl.nn.hr_skip module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.mlp.html">cyto_dl.nn.mlp module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.res_unit.html">cyto_dl.nn.res_unit module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.spatial_transformer.html">cyto_dl.nn.spatial_transformer module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
+</ul>
+</li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
+</ul>
+</li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.utils.array.html">cyto_dl.utils.array module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.utils.checkpoint.html">cyto_dl.utils.checkpoint module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.utils.config.html">cyto_dl.utils.config module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.utils.download_test_data.html">cyto_dl.utils.download_test_data module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.utils.dummy_dataset.html">cyto_dl.utils.dummy_dataset module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.utils.pylogger.html">cyto_dl.utils.pylogger module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.utils.rich_utils.html">cyto_dl.utils.rich_utils module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.utils.rotation.html">cyto_dl.utils.rotation module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.utils.template_utils.html">cyto_dl.utils.template_utils module</a></li>
+</ul>
+</li>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.compile.html">cyto_dl.compile module</a></li>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.eval.html">cyto_dl.eval module</a></li>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.train.html">cyto_dl.train module</a></li>
+</ul>
+</li>
+</ul>
+</li>
+</ul>
+
+</div>
+</div>
+
+      </div>
+      
+    </div>
+  </aside>
+  <div class="main">
+    <div class="content">
+      <div class="article-container">
+        <a href="#" class="back-to-top muted-link">
+          <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24">
+            <path d="M13 20h-2V8l-5.5 5.5-1.42-1.42L12 4.16l7.92 7.92-1.42 1.42L13 8v12z"></path>
+          </svg>
+          <span>Back to top</span>
+        </a>
+        <div class="content-icon-container">
+          
+<div class="theme-toggle-container theme-toggle-content">
+            <button class="theme-toggle">
+              <div class="visually-hidden">Toggle Light / Dark / Auto color theme</div>
+              <svg class="theme-icon-when-auto"><use href="#svg-sun-half"></use></svg>
+              <svg class="theme-icon-when-dark"><use href="#svg-moon"></use></svg>
+              <svg class="theme-icon-when-light"><use href="#svg-sun"></use></svg>
+            </button>
+          </div>
+          <label class="toc-overlay-icon toc-content-icon" for="__toc">
+            <div class="visually-hidden">Toggle table of contents sidebar</div>
+            <i class="icon"><svg><use href="#svg-toc"></use></svg></i>
+          </label>
+        </div>
+        <article role="main">
+          <section id="module-cyto_dl.nn.vits.decoder">
+<span id="cyto-dl-nn-vits-decoder-module"></span><h1>cyto_dl.nn.vits.decoder module<a class="headerlink" href="#module-cyto_dl.nn.vits.decoder" title="Permalink to this heading">#</a></h1>
+<dl class="py class">
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.decoder.CrossMAE_Decoder">
+<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">cyto_dl.nn.vits.decoder.</span></span><span class="sig-name descname"><span class="pre">CrossMAE_Decoder</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">num_patches</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">]</span></span></span></em>, <em class="sig-param"><span class="n"><span class="pre">spatial_dims</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">3</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">patch_size</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">]</span></span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">4</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">enc_dim</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">768</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">emb_dim</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">192</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">num_layer</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">4</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">num_head</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">3</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">has_cls_token</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">bool</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">True</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">learnable_pos_embedding</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">bool</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">True</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/decoder.html#CrossMAE_Decoder"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.decoder.CrossMAE_Decoder" title="Permalink to this definition">#</a></dt>
+<dd><p>Bases: <a class="reference internal" href="#cyto_dl.nn.vits.decoder.MAE_Decoder" title="cyto_dl.nn.vits.decoder.MAE_Decoder"><code class="xref py py-class docutils literal notranslate"><span class="pre">MAE_Decoder</span></code></a></p>
+<p>Decoder inspired by [CrossMAE](<a class="reference external" href="https://crossmae.github.io/">https://crossmae.github.io/</a>) where masked tokens only attend
+to visible tokens.</p>
+<dl class="field-list simple">
+<dt class="field-odd">Parameters<span class="colon">:</span></dt>
+<dd class="field-odd"><ul class="simple">
+<li><p><strong>num_patches</strong> (<em>List[int], int</em>) – Number of patches in each dimension. If int, the same number of patches is used for all dimensions.</p></li>
+<li><p><strong>patch_size</strong> (<em>Tuple[int]</em>) – Size of each patch in each dimension. If int, the same patch size is used for all dimensions.</p></li>
+<li><p><strong>enc_dim</strong> (<em>int</em>) – Dimension of encoder</p></li>
+<li><p><strong>emb_dim</strong> (<em>int</em>) – Dimension of embedding</p></li>
+<li><p><strong>num_layer</strong> (<em>int</em>) – Number of transformer layers</p></li>
+<li><p><strong>num_head</strong> (<em>int</em>) – Number of heads in transformer</p></li>
+<li><p><strong>has_cls_token</strong> (<em>bool</em>) – Whether encoder features have a cls token</p></li>
+<li><p><strong>learnable_pos_embedding</strong> (<em>bool</em>) – If True, learnable positional embeddings are used. If False, fixed sin/cos positional embeddings are used. Empirically, fixed positional embeddings work better for brightfield images.</p></li>
+</ul>
+</dd>
+</dl>
+<dl class="py method">
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.decoder.CrossMAE_Decoder.forward">
+<span class="sig-name descname"><span class="pre">forward</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">features</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">forward_indexes</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">backward_indexes</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/decoder.html#CrossMAE_Decoder.forward"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.decoder.CrossMAE_Decoder.forward" title="Permalink to this definition">#</a></dt>
+<dd></dd></dl>
+
+</dd></dl>
+
+<dl class="py class">
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.decoder.MAE_Decoder">
+<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">cyto_dl.nn.vits.decoder.</span></span><span class="sig-name descname"><span class="pre">MAE_Decoder</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">num_patches</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">]</span></span></span></em>, <em class="sig-param"><span class="n"><span class="pre">spatial_dims</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">3</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">patch_size</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">]</span></span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">4</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">enc_dim</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">768</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">emb_dim</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">192</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">num_layer</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">4</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">num_head</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">3</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">has_cls_token</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">bool</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">False</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">learnable_pos_embedding</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">bool</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">True</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/decoder.html#MAE_Decoder"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.decoder.MAE_Decoder" title="Permalink to this definition">#</a></dt>
+<dd><p>Bases: <code class="xref py py-class docutils literal notranslate"><span class="pre">Module</span></code></p>
+<dl class="field-list simple">
+<dt class="field-odd">Parameters<span class="colon">:</span></dt>
+<dd class="field-odd"><ul class="simple">
+<li><p><strong>num_patches</strong> (<em>List[int], int</em>) – Number of patches in each dimension. If int, the same number of patches is used for all dimensions.</p></li>
+<li><p><strong>patch_size</strong> (<em>Tuple[int], int</em>) – Size of each patch. If int, the same patch size is used for all dimensions.</p></li>
+<li><p><strong>enc_dim</strong> (<em>int</em>) – Dimension of encoder</p></li>
+<li><p><strong>emb_dim</strong> (<em>int</em>) – Dimension of decoder</p></li>
+<li><p><strong>num_layer</strong> (<em>int</em>) – Number of transformer layers</p></li>
+<li><p><strong>num_head</strong> (<em>int</em>) – Number of heads in transformer</p></li>
+<li><p><strong>has_cls_token</strong> (<em>bool</em>) – Whether encoder features have a cls token</p></li>
+<li><p><strong>learnable_pos_embedding</strong> (<em>bool</em>) – If True, learnable positional embeddings are used. If False, fixed sin/cos positional embeddings. Empirically, fixed positional embeddings work better for brightfield images.</p></li>
+</ul>
+</dd>
+</dl>
+<dl class="py method">
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.decoder.MAE_Decoder.add_mask_tokens">
+<span class="sig-name descname"><span class="pre">add_mask_tokens</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">features</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">backward_indexes</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/decoder.html#MAE_Decoder.add_mask_tokens"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.decoder.MAE_Decoder.add_mask_tokens" title="Permalink to this definition">#</a></dt>
+<dd></dd></dl>
+
+<dl class="py method">
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.decoder.MAE_Decoder.adjust_indices_for_cls">
+<span class="sig-name descname"><span class="pre">adjust_indices_for_cls</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">indexes</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/decoder.html#MAE_Decoder.adjust_indices_for_cls"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.decoder.MAE_Decoder.adjust_indices_for_cls" title="Permalink to this definition">#</a></dt>
+<dd></dd></dl>
+
+<dl class="py method">
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.decoder.MAE_Decoder.forward">
+<span class="sig-name descname"><span class="pre">forward</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">features</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">forward_indexes</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">backward_indexes</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/decoder.html#MAE_Decoder.forward"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.decoder.MAE_Decoder.forward" title="Permalink to this definition">#</a></dt>
+<dd></dd></dl>
+
+<dl class="py method">
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.decoder.MAE_Decoder.init_weight">
+<span class="sig-name descname"><span class="pre">init_weight</span></span><span class="sig-paren">(</span><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/decoder.html#MAE_Decoder.init_weight"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.decoder.MAE_Decoder.init_weight" title="Permalink to this definition">#</a></dt>
+<dd></dd></dl>
+
+</dd></dl>
+
+</section>
+
+        </article>
+      </div>
+      <footer>
+        
+        <div class="related-pages">
+          <a class="next-page" href="cyto_dl.nn.vits.encoder.html">
+              <div class="page-info">
+                <div class="context">
+                  <span>Next</span>
+                </div>
+                <div class="title">cyto_dl.nn.vits.encoder module</div>
+              </div>
+              <svg class="furo-related-icon"><use href="#svg-arrow-right"></use></svg>
+            </a>
+          <a class="prev-page" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">
+              <svg class="furo-related-icon"><use href="#svg-arrow-right"></use></svg>
+              <div class="page-info">
+                <div class="context">
+                  <span>Previous</span>
+                </div>
+                
+                <div class="title">cyto_dl.nn.vits.blocks.masked_unit_attention module</div>
+                
+              </div>
+            </a>
+        </div>
+        <div class="bottom-of-page">
+          <div class="left-details">
+            <div class="copyright">
+                Copyright &#169; 2023
+            </div>
+            Made with <a href="https://www.sphinx-doc.org/">Sphinx</a> and <a class="muted-link" href="https://pradyunsg.me">@pradyunsg</a>'s
+            
+            <a href="https://github.com/pradyunsg/furo">Furo</a>
+            
+          </div>
+          <div class="right-details">
+            
+          </div>
+        </div>
+        
+      </footer>
+    </div>
+    <aside class="toc-drawer">
+      
+      
+      <div class="toc-sticky toc-scroll">
+        <div class="toc-title-container">
+          <span class="toc-title">
+            On this page
+          </span>
+        </div>
+        <div class="toc-tree-container">
+          <div class="toc-tree">
+            <ul>
+<li><a class="reference internal" href="#">cyto_dl.nn.vits.decoder module</a><ul>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.decoder.CrossMAE_Decoder"><code class="docutils literal notranslate"><span class="pre">CrossMAE_Decoder</span></code></a><ul>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.decoder.CrossMAE_Decoder.forward"><code class="docutils literal notranslate"><span class="pre">CrossMAE_Decoder.forward()</span></code></a></li>
+</ul>
+</li>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.decoder.MAE_Decoder"><code class="docutils literal notranslate"><span class="pre">MAE_Decoder</span></code></a><ul>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.decoder.MAE_Decoder.add_mask_tokens"><code class="docutils literal notranslate"><span class="pre">MAE_Decoder.add_mask_tokens()</span></code></a></li>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.decoder.MAE_Decoder.adjust_indices_for_cls"><code class="docutils literal notranslate"><span class="pre">MAE_Decoder.adjust_indices_for_cls()</span></code></a></li>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.decoder.MAE_Decoder.forward"><code class="docutils literal notranslate"><span class="pre">MAE_Decoder.forward()</span></code></a></li>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.decoder.MAE_Decoder.init_weight"><code class="docutils literal notranslate"><span class="pre">MAE_Decoder.init_weight()</span></code></a></li>
+</ul>
+</li>
+</ul>
+</li>
+</ul>
+
+          </div>
+        </div>
+      </div>
+      
+      
+    </aside>
+  </div>
+</div><script data-url_root="./" id="documentation_options" src="_static/documentation_options.js"></script>
+    <script src="_static/doctools.js"></script>
+    <script src="_static/sphinx_highlight.js"></script>
+    <script src="_static/scripts/furo.js"></script>
+    </body>
+</html>
\ No newline at end of file
diff --git a/cyto_dl.nn.vits.encoder.html b/cyto_dl.nn.vits.encoder.html
new file mode 100644
index 000000000..0ad201af2
--- /dev/null
+++ b/cyto_dl.nn.vits.encoder.html
@@ -0,0 +1,656 @@
+<!doctype html>
+<html class="no-js" lang="en">
+  <head><meta charset="utf-8"/>
+    <meta name="viewport" content="width=device-width,initial-scale=1"/>
+    <meta name="color-scheme" content="light dark"><meta name="generator" content="Docutils 0.19: https://docutils.sourceforge.io/" />
+<link rel="index" title="Index" href="genindex.html" /><link rel="search" title="Search" href="search.html" /><link rel="next" title="cyto_dl.nn.vits.mae module" href="cyto_dl.nn.vits.mae.html" /><link rel="prev" title="cyto_dl.nn.vits.decoder module" href="cyto_dl.nn.vits.decoder.html" />
+
+    <!-- Generated with Sphinx 6.1.3 and Furo 2023.03.23 -->
+        <title>cyto_dl.nn.vits.encoder module - cyto-dl 0.4.0 documentation</title>
+      <link rel="stylesheet" type="text/css" href="_static/pygments.css" />
+    <link rel="stylesheet" type="text/css" href="_static/styles/furo.css?digest=fad236701ea90a88636c2a8c73b44ae642ed2a53" />
+    <link rel="stylesheet" type="text/css" href="_static/styles/furo-extensions.css?digest=30d1aed668e5c3a91c3e3bf6a60b675221979f0e" />
+    
+    
+
+
+<style>
+  body {
+    --color-code-background: #eeffcc;
+  --color-code-foreground: black;
+  
+  }
+  @media not print {
+    body[data-theme="dark"] {
+      --color-code-background: #202020;
+  --color-code-foreground: #d0d0d0;
+  
+    }
+    @media (prefers-color-scheme: dark) {
+      body:not([data-theme="light"]) {
+        --color-code-background: #202020;
+  --color-code-foreground: #d0d0d0;
+  
+      }
+    }
+  }
+</style></head>
+  <body>
+    
+    <script>
+      document.body.dataset.theme = localStorage.getItem("theme") || "auto";
+    </script>
+    
+
+<svg xmlns="http://www.w3.org/2000/svg" style="display: none;">
+  <symbol id="svg-toc" viewBox="0 0 24 24">
+    <title>Contents</title>
+    <svg stroke="currentColor" fill="currentColor" stroke-width="0" viewBox="0 0 1024 1024">
+      <path d="M408 442h480c4.4 0 8-3.6 8-8v-56c0-4.4-3.6-8-8-8H408c-4.4 0-8 3.6-8 8v56c0 4.4 3.6 8 8 8zm-8 204c0 4.4 3.6 8 8 8h480c4.4 0 8-3.6 8-8v-56c0-4.4-3.6-8-8-8H408c-4.4 0-8 3.6-8 8v56zm504-486H120c-4.4 0-8 3.6-8 8v56c0 4.4 3.6 8 8 8h784c4.4 0 8-3.6 8-8v-56c0-4.4-3.6-8-8-8zm0 632H120c-4.4 0-8 3.6-8 8v56c0 4.4 3.6 8 8 8h784c4.4 0 8-3.6 8-8v-56c0-4.4-3.6-8-8-8zM115.4 518.9L271.7 642c5.8 4.6 14.4.5 14.4-6.9V388.9c0-7.4-8.5-11.5-14.4-6.9L115.4 505.1a8.74 8.74 0 0 0 0 13.8z"/>
+    </svg>
+  </symbol>
+  <symbol id="svg-menu" viewBox="0 0 24 24">
+    <title>Menu</title>
+    <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
+      stroke-width="2" stroke-linecap="round" stroke-linejoin="round" class="feather-menu">
+      <line x1="3" y1="12" x2="21" y2="12"></line>
+      <line x1="3" y1="6" x2="21" y2="6"></line>
+      <line x1="3" y1="18" x2="21" y2="18"></line>
+    </svg>
+  </symbol>
+  <symbol id="svg-arrow-right" viewBox="0 0 24 24">
+    <title>Expand</title>
+    <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
+      stroke-width="2" stroke-linecap="round" stroke-linejoin="round" class="feather-chevron-right">
+      <polyline points="9 18 15 12 9 6"></polyline>
+    </svg>
+  </symbol>
+  <symbol id="svg-sun" viewBox="0 0 24 24">
+    <title>Light mode</title>
+    <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
+      stroke-width="1.5" stroke-linecap="round" stroke-linejoin="round" class="feather-sun">
+      <circle cx="12" cy="12" r="5"></circle>
+      <line x1="12" y1="1" x2="12" y2="3"></line>
+      <line x1="12" y1="21" x2="12" y2="23"></line>
+      <line x1="4.22" y1="4.22" x2="5.64" y2="5.64"></line>
+      <line x1="18.36" y1="18.36" x2="19.78" y2="19.78"></line>
+      <line x1="1" y1="12" x2="3" y2="12"></line>
+      <line x1="21" y1="12" x2="23" y2="12"></line>
+      <line x1="4.22" y1="19.78" x2="5.64" y2="18.36"></line>
+      <line x1="18.36" y1="5.64" x2="19.78" y2="4.22"></line>
+    </svg>
+  </symbol>
+  <symbol id="svg-moon" viewBox="0 0 24 24">
+    <title>Dark mode</title>
+    <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
+      stroke-width="1.5" stroke-linecap="round" stroke-linejoin="round" class="icon-tabler-moon">
+      <path stroke="none" d="M0 0h24v24H0z" fill="none" />
+      <path d="M12 3c.132 0 .263 0 .393 0a7.5 7.5 0 0 0 7.92 12.446a9 9 0 1 1 -8.313 -12.454z" />
+    </svg>
+  </symbol>
+  <symbol id="svg-sun-half" viewBox="0 0 24 24">
+    <title>Auto light/dark mode</title>
+    <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
+      stroke-width="1.5" stroke-linecap="round" stroke-linejoin="round" class="icon-tabler-shadow">
+      <path stroke="none" d="M0 0h24v24H0z" fill="none"/>
+      <circle cx="12" cy="12" r="9" />
+      <path d="M13 12h5" />
+      <path d="M13 15h4" />
+      <path d="M13 18h1" />
+      <path d="M13 9h4" />
+      <path d="M13 6h1" />
+    </svg>
+  </symbol>
+</svg>
+
+<input type="checkbox" class="sidebar-toggle" name="__navigation" id="__navigation">
+<input type="checkbox" class="sidebar-toggle" name="__toc" id="__toc">
+<label class="overlay sidebar-overlay" for="__navigation">
+  <div class="visually-hidden">Hide navigation sidebar</div>
+</label>
+<label class="overlay toc-overlay" for="__toc">
+  <div class="visually-hidden">Hide table of contents sidebar</div>
+</label>
+
+
+
+<div class="page">
+  <header class="mobile-header">
+    <div class="header-left">
+      <label class="nav-overlay-icon" for="__navigation">
+        <div class="visually-hidden">Toggle site navigation sidebar</div>
+        <i class="icon"><svg><use href="#svg-menu"></use></svg></i>
+      </label>
+    </div>
+    <div class="header-center">
+      <a href="index.html"><div class="brand">cyto-dl 0.4.0 documentation</div></a>
+    </div>
+    <div class="header-right">
+      <div class="theme-toggle-container theme-toggle-header">
+        <button class="theme-toggle">
+          <div class="visually-hidden">Toggle Light / Dark / Auto color theme</div>
+          <svg class="theme-icon-when-auto"><use href="#svg-sun-half"></use></svg>
+          <svg class="theme-icon-when-dark"><use href="#svg-moon"></use></svg>
+          <svg class="theme-icon-when-light"><use href="#svg-sun"></use></svg>
+        </button>
+      </div>
+      <label class="toc-overlay-icon toc-header-icon" for="__toc">
+        <div class="visually-hidden">Toggle table of contents sidebar</div>
+        <i class="icon"><svg><use href="#svg-toc"></use></svg></i>
+      </label>
+    </div>
+  </header>
+  <aside class="sidebar-drawer">
+    <div class="sidebar-container">
+      
+      <div class="sidebar-sticky"><a class="sidebar-brand" href="index.html">
+  
+  
+  <span class="sidebar-brand-text">cyto-dl 0.4.0 documentation</span>
+  
+</a><form class="sidebar-search-container" method="get" action="search.html" role="search">
+  <input class="sidebar-search" placeholder="Search" name="q" aria-label="Search">
+  <input type="hidden" name="check_keywords" value="yes">
+  <input type="hidden" name="area" value="default">
+</form>
+<div id="searchbox"></div><div class="sidebar-scroll"><div class="sidebar-tree">
+  <p class="caption" role="heading"><span class="caption-text">Contents:</span></p>
+<ul class="current">
+<li class="toctree-l1"><a class="reference internal" href="index.html">Home</a></li>
+<li class="toctree-l1"><a class="reference internal" href="installation.html">Installation</a></li>
+<li class="toctree-l1"><a class="reference internal" href="using_examples.html">Modifying the Example Configs</a></li>
+<li class="toctree-l1"><a class="reference internal" href="api.html">Using the API</a></li>
+<li class="toctree-l1 current has-children"><a class="reference internal" href="modules.html">Package modules</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-1" name="toctree-checkbox-1" role="switch" type="checkbox"/><label for="toctree-checkbox-1"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
+<li class="toctree-l2 current has-children"><a class="reference internal" href="cyto_dl.html">cyto_dl package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-2" name="toctree-checkbox-2" role="switch" type="checkbox"/><label for="toctree-checkbox-2"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.api.html">cyto_dl.api package</a><input class="toctree-checkbox" id="toctree-checkbox-3" name="toctree-checkbox-3" role="switch" type="checkbox"/><label for="toctree-checkbox-3"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.api.cyto_dl_model.html">cyto_dl.api.cyto_dl_model package</a><input class="toctree-checkbox" id="toctree-checkbox-4" name="toctree-checkbox-4" role="switch" type="checkbox"/><label for="toctree-checkbox-4"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.api.cyto_dl_model.cyto_dl_base_model.html">cyto_dl.api.cyto_dl_model.cyto_dl_base_model module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.api.cyto_dl_model.segmentation_plugin_model.html">cyto_dl.api.cyto_dl_model.segmentation_plugin_model module</a></li>
+</ul>
+</li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.api.data.html">cyto_dl.api.data module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.api.model.html">cyto_dl.api.model module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.callbacks.html">cyto_dl.callbacks package</a><input class="toctree-checkbox" id="toctree-checkbox-5" name="toctree-checkbox-5" role="switch" type="checkbox"/><label for="toctree-checkbox-5"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.callbacks.callback_utils.html">cyto_dl.callbacks.callback_utils module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.callbacks.csv_saver.html">cyto_dl.callbacks.csv_saver module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.callbacks.image_saver.html">cyto_dl.callbacks.image_saver module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.callbacks.latent_walk.html">cyto_dl.callbacks.latent_walk module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.callbacks.layer_freeze.html">cyto_dl.callbacks.layer_freeze module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.callbacks.model_utils.html">cyto_dl.callbacks.model_utils module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.callbacks.outlier_detection.html">cyto_dl.callbacks.outlier_detection module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.dataframe.html">cyto_dl.dataframe package</a><input class="toctree-checkbox" id="toctree-checkbox-6" name="toctree-checkbox-6" role="switch" type="checkbox"/><label for="toctree-checkbox-6"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.dataframe.transforms.html">cyto_dl.dataframe.transforms package</a><input class="toctree-checkbox" id="toctree-checkbox-7" name="toctree-checkbox-7" role="switch" type="checkbox"/><label for="toctree-checkbox-7"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.dataframe.transforms.filter.html">cyto_dl.dataframe.transforms.filter module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.dataframe.transforms.group_cols.html">cyto_dl.dataframe.transforms.group_cols module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.dataframe.transforms.misc.html">cyto_dl.dataframe.transforms.misc module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.dataframe.transforms.split.html">cyto_dl.dataframe.transforms.split module</a></li>
+</ul>
+</li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.dataframe.readers.html">cyto_dl.dataframe.readers module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.datamodules.html">cyto_dl.datamodules package</a><input class="toctree-checkbox" id="toctree-checkbox-8" name="toctree-checkbox-8" role="switch" type="checkbox"/><label for="toctree-checkbox-8"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.datamodules.dataframe.html">cyto_dl.datamodules.dataframe package</a><input class="toctree-checkbox" id="toctree-checkbox-9" name="toctree-checkbox-9" role="switch" type="checkbox"/><label for="toctree-checkbox-9"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.datamodules.dataframe.dataframe_datamodule.html">cyto_dl.datamodules.dataframe.dataframe_datamodule module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.datamodules.dataframe.grouped_dataframe_datamodule.html">cyto_dl.datamodules.dataframe.grouped_dataframe_datamodule module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.datamodules.dataframe.utils.html">cyto_dl.datamodules.dataframe.utils module</a></li>
+</ul>
+</li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.datamodules.array.html">cyto_dl.datamodules.array module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.datamodules.data_dict.html">cyto_dl.datamodules.data_dict module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.datamodules.folder.html">cyto_dl.datamodules.folder module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.datamodules.multidim_image.html">cyto_dl.datamodules.multidim_image module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.datamodules.smartcache.html">cyto_dl.datamodules.smartcache module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.datamodules.torchvision.html">cyto_dl.datamodules.torchvision module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.image.html">cyto_dl.image package</a><input class="toctree-checkbox" id="toctree-checkbox-10" name="toctree-checkbox-10" role="switch" type="checkbox"/><label for="toctree-checkbox-10"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.image.io.html">cyto_dl.image.io package</a><input class="toctree-checkbox" id="toctree-checkbox-11" name="toctree-checkbox-11" role="switch" type="checkbox"/><label for="toctree-checkbox-11"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.io.aicsimage_loader.html">cyto_dl.image.io.aicsimage_loader module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.io.monai_bio_reader.html">cyto_dl.image.io.monai_bio_reader module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.io.numpy_reader.html">cyto_dl.image.io.numpy_reader module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.io.ome_zarr_reader.html">cyto_dl.image.io.ome_zarr_reader module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.io.polygon_loader.html">cyto_dl.image.io.polygon_loader module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.io.skimage_reader.html">cyto_dl.image.io.skimage_reader module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.image.transforms.html">cyto_dl.image.transforms package</a><input class="toctree-checkbox" id="toctree-checkbox-12" name="toctree-checkbox-12" role="switch" type="checkbox"/><label for="toctree-checkbox-12"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.add_meta.html">cyto_dl.image.transforms.add_meta module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.bright_sampler.html">cyto_dl.image.transforms.bright_sampler module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.clip.html">cyto_dl.image.transforms.clip module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.contrastadjust.html">cyto_dl.image.transforms.contrastadjust module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.generate_jepa_masks.html">cyto_dl.image.transforms.generate_jepa_masks module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.merge.html">cyto_dl.image.transforms.merge module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.multiscale_cropper.html">cyto_dl.image.transforms.multiscale_cropper module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.pad.html">cyto_dl.image.transforms.pad module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.project.html">cyto_dl.image.transforms.project module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.rotation_mask_transform.html">cyto_dl.image.transforms.rotation_mask_transform module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.save.html">cyto_dl.image.transforms.save module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.so2_random_rotation.html">cyto_dl.image.transforms.so2_random_rotation module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.image.transforms.track_transforms.html">cyto_dl.image.transforms.track_transforms module</a></li>
+</ul>
+</li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.loggers.html">cyto_dl.loggers package</a><input class="toctree-checkbox" id="toctree-checkbox-13" name="toctree-checkbox-13" role="switch" type="checkbox"/><label for="toctree-checkbox-13"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.loggers.mlflow.html">cyto_dl.loggers.mlflow module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.models.html">cyto_dl.models package</a><input class="toctree-checkbox" id="toctree-checkbox-14" name="toctree-checkbox-14" role="switch" type="checkbox"/><label for="toctree-checkbox-14"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.models.classification.html">cyto_dl.models.classification package</a><input class="toctree-checkbox" id="toctree-checkbox-15" name="toctree-checkbox-15" role="switch" type="checkbox"/><label for="toctree-checkbox-15"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.classification.classification.html">cyto_dl.models.classification.classification module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.classification.timepoint_classification.html">cyto_dl.models.classification.timepoint_classification module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.models.contrastive.html">cyto_dl.models.contrastive package</a><input class="toctree-checkbox" id="toctree-checkbox-16" name="toctree-checkbox-16" role="switch" type="checkbox"/><label for="toctree-checkbox-16"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.contrastive.contrastive.html">cyto_dl.models.contrastive.contrastive module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.models.handlers.html">cyto_dl.models.handlers package</a><input class="toctree-checkbox" id="toctree-checkbox-17" name="toctree-checkbox-17" role="switch" type="checkbox"/><label for="toctree-checkbox-17"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.handlers.base_handler.html">cyto_dl.models.handlers.base_handler module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.handlers.image_handler.html">cyto_dl.models.handlers.image_handler module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.handlers.load_image_patch.html">cyto_dl.models.handlers.load_image_patch module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.models.im2im.html">cyto_dl.models.im2im package</a><input class="toctree-checkbox" id="toctree-checkbox-18" name="toctree-checkbox-18" role="switch" type="checkbox"/><label for="toctree-checkbox-18"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.models.im2im.utils.html">cyto_dl.models.im2im.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-19" name="toctree-checkbox-19" role="switch" type="checkbox"/><label for="toctree-checkbox-19"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.models.im2im.utils.postprocessing.html">cyto_dl.models.im2im.utils.postprocessing package</a><input class="toctree-checkbox" id="toctree-checkbox-20" name="toctree-checkbox-20" role="switch" type="checkbox"/><label for="toctree-checkbox-20"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.models.im2im.utils.postprocessing.act_thresh_label.html">cyto_dl.models.im2im.utils.postprocessing.act_thresh_label module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.models.im2im.utils.postprocessing.auto_thresh.html">cyto_dl.models.im2im.utils.postprocessing.auto_thresh module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.models.im2im.utils.postprocessing.dict_to_im.html">cyto_dl.models.im2im.utils.postprocessing.dict_to_im module</a></li>
+</ul>
+</li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.models.im2im.utils.instance_seg.html">cyto_dl.models.im2im.utils.instance_seg module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.models.im2im.utils.noise_annealer.html">cyto_dl.models.im2im.utils.noise_annealer module</a></li>
+</ul>
+</li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.im2im.gan.html">cyto_dl.models.im2im.gan module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.im2im.multi_task.html">cyto_dl.models.im2im.multi_task module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.models.jepa.html">cyto_dl.models.jepa package</a><input class="toctree-checkbox" id="toctree-checkbox-21" name="toctree-checkbox-21" role="switch" type="checkbox"/><label for="toctree-checkbox-21"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.jepa.ijepa.html">cyto_dl.models.jepa.ijepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.jepa.iwm.html">cyto_dl.models.jepa.iwm module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.jepa.jepa_base.html">cyto_dl.models.jepa.jepa_base module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.models.utils.html">cyto_dl.models.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-22" name="toctree-checkbox-22" role="switch" type="checkbox"/><label for="toctree-checkbox-22"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.utils.mlflow.html">cyto_dl.models.utils.mlflow module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.utils.utils.html">cyto_dl.models.utils.utils module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.models.vae.html">cyto_dl.models.vae package</a><input class="toctree-checkbox" id="toctree-checkbox-23" name="toctree-checkbox-23" role="switch" type="checkbox"/><label for="toctree-checkbox-23"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.models.vae.o2_spharm_vae.html">cyto_dl.models.vae.o2_spharm_vae package</a><input class="toctree-checkbox" id="toctree-checkbox-24" name="toctree-checkbox-24" role="switch" type="checkbox"/><label for="toctree-checkbox-24"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.models.vae.o2_spharm_vae.o2_spharm_encoder.html">cyto_dl.models.vae.o2_spharm_vae.o2_spharm_encoder module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.models.vae.o2_spharm_vae.o2_spharm_vae.html">cyto_dl.models.vae.o2_spharm_vae.o2_spharm_vae module</a></li>
+</ul>
+</li>
+<li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.models.vae.priors.html">cyto_dl.models.vae.priors package</a><input class="toctree-checkbox" id="toctree-checkbox-25" name="toctree-checkbox-25" role="switch" type="checkbox"/><label for="toctree-checkbox-25"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.models.vae.priors.abstract_prior.html">cyto_dl.models.vae.priors.abstract_prior module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.models.vae.priors.gaussian.html">cyto_dl.models.vae.priors.gaussian module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.models.vae.priors.identity_prior.html">cyto_dl.models.vae.priors.identity_prior module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.models.vae.priors.joint_prior.html">cyto_dl.models.vae.priors.joint_prior module</a></li>
+</ul>
+</li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.vae.base_vae.html">cyto_dl.models.vae.base_vae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.vae.image_canon_vae.html">cyto_dl.models.vae.image_canon_vae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.vae.image_encoder.html">cyto_dl.models.vae.image_encoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.vae.image_vae.html">cyto_dl.models.vae.image_vae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.vae.latent_loss_vae.html">cyto_dl.models.vae.latent_loss_vae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.vae.point_cloud_vae.html">cyto_dl.models.vae.point_cloud_vae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.models.vae.tabular_vae.html">cyto_dl.models.vae.tabular_vae module</a></li>
+</ul>
+</li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.models.base_model.html">cyto_dl.models.base_model module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.models.basic_model.html">cyto_dl.models.basic_model module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 current has-children"><a class="reference internal" href="cyto_dl.nn.html">cyto_dl.nn package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-26" name="toctree-checkbox-26" role="switch" type="checkbox"/><label for="toctree-checkbox-26"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.discriminators.html">cyto_dl.nn.discriminators package</a><input class="toctree-checkbox" id="toctree-checkbox-27" name="toctree-checkbox-27" role="switch" type="checkbox"/><label for="toctree-checkbox-27"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.discriminators.multi_scale_discriminator.html">cyto_dl.nn.discriminators.multi_scale_discriminator module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.discriminators.n_layer_discriminator.html">cyto_dl.nn.discriminators.n_layer_discriminator module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.head.html">cyto_dl.nn.head package</a><input class="toctree-checkbox" id="toctree-checkbox-28" name="toctree-checkbox-28" role="switch" type="checkbox"/><label for="toctree-checkbox-28"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.head.base_head.html">cyto_dl.nn.head.base_head module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.head.gan_head.html">cyto_dl.nn.head.gan_head module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.head.gan_head_superres.html">cyto_dl.nn.head.gan_head_superres module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.head.mae_head.html">cyto_dl.nn.head.mae_head module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.head.mask_head.html">cyto_dl.nn.head.mask_head module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.head.res_blocks_head.html">cyto_dl.nn.head.res_blocks_head module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.head.vic_reg.html">cyto_dl.nn.head.vic_reg module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.losses.html">cyto_dl.nn.losses package</a><input class="toctree-checkbox" id="toctree-checkbox-29" name="toctree-checkbox-29" role="switch" type="checkbox"/><label for="toctree-checkbox-29"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.adversarial_loss.html">cyto_dl.nn.losses.adversarial_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.chamfer_loss.html">cyto_dl.nn.losses.chamfer_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.continuous_bernoulli.html">cyto_dl.nn.losses.continuous_bernoulli module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.cosine_loss.html">cyto_dl.nn.losses.cosine_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.gan_loss.html">cyto_dl.nn.losses.gan_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.gaussian_nll_loss.html">cyto_dl.nn.losses.gaussian_nll_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.geomloss.html">cyto_dl.nn.losses.geomloss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.loss_wrapper.html">cyto_dl.nn.losses.loss_wrapper module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.spharm_loss.html">cyto_dl.nn.losses.spharm_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.threshold_loss.html">cyto_dl.nn.losses.threshold_loss module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.vic_reg.html">cyto_dl.nn.losses.vic_reg module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.weibull.html">cyto_dl.nn.losses.weibull module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.losses.weighted_mse_loss.html">cyto_dl.nn.losses.weighted_mse_loss module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.point_cloud.html">cyto_dl.nn.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-30" name="toctree-checkbox-30" role="switch" type="checkbox"/><label for="toctree-checkbox-30"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.point_cloud.dgcnn.html">cyto_dl.nn.point_cloud.dgcnn module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.point_cloud.folding_net.html">cyto_dl.nn.point_cloud.folding_net module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.point_cloud.graph_functions.html">cyto_dl.nn.point_cloud.graph_functions module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.point_cloud.vnn.html">cyto_dl.nn.point_cloud.vnn module</a></li>
+</ul>
+</li>
+<li class="toctree-l4 current has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
+<li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
+</ul>
+</li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5 current current-page"><a class="current reference internal" href="#">cyto_dl.nn.vits.encoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
+</ul>
+</li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.hr_skip.html">cyto_dl.nn.hr_skip module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.mlp.html">cyto_dl.nn.mlp module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.res_unit.html">cyto_dl.nn.res_unit module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.spatial_transformer.html">cyto_dl.nn.spatial_transformer module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
+</ul>
+</li>
+</ul>
+</li>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
+</ul>
+</li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.utils.array.html">cyto_dl.utils.array module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.utils.checkpoint.html">cyto_dl.utils.checkpoint module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.utils.config.html">cyto_dl.utils.config module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.utils.download_test_data.html">cyto_dl.utils.download_test_data module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.utils.dummy_dataset.html">cyto_dl.utils.dummy_dataset module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.utils.pylogger.html">cyto_dl.utils.pylogger module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.utils.rich_utils.html">cyto_dl.utils.rich_utils module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.utils.rotation.html">cyto_dl.utils.rotation module</a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.utils.template_utils.html">cyto_dl.utils.template_utils module</a></li>
+</ul>
+</li>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.compile.html">cyto_dl.compile module</a></li>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.eval.html">cyto_dl.eval module</a></li>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.train.html">cyto_dl.train module</a></li>
+</ul>
+</li>
+</ul>
+</li>
+</ul>
+
+</div>
+</div>
+
+      </div>
+      
+    </div>
+  </aside>
+  <div class="main">
+    <div class="content">
+      <div class="article-container">
+        <a href="#" class="back-to-top muted-link">
+          <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24">
+            <path d="M13 20h-2V8l-5.5 5.5-1.42-1.42L12 4.16l7.92 7.92-1.42 1.42L13 8v12z"></path>
+          </svg>
+          <span>Back to top</span>
+        </a>
+        <div class="content-icon-container">
+          
+<div class="theme-toggle-container theme-toggle-content">
+            <button class="theme-toggle">
+              <div class="visually-hidden">Toggle Light / Dark / Auto color theme</div>
+              <svg class="theme-icon-when-auto"><use href="#svg-sun-half"></use></svg>
+              <svg class="theme-icon-when-dark"><use href="#svg-moon"></use></svg>
+              <svg class="theme-icon-when-light"><use href="#svg-sun"></use></svg>
+            </button>
+          </div>
+          <label class="toc-overlay-icon toc-content-icon" for="__toc">
+            <div class="visually-hidden">Toggle table of contents sidebar</div>
+            <i class="icon"><svg><use href="#svg-toc"></use></svg></i>
+          </label>
+        </div>
+        <article role="main">
+          <section id="module-cyto_dl.nn.vits.encoder">
+<span id="cyto-dl-nn-vits-encoder-module"></span><h1>cyto_dl.nn.vits.encoder module<a class="headerlink" href="#module-cyto_dl.nn.vits.encoder" title="Permalink to this heading">#</a></h1>
+<dl class="py class">
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.encoder.HieraEncoder">
+<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">cyto_dl.nn.vits.encoder.</span></span><span class="sig-name descname"><span class="pre">HieraEncoder</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">num_patches</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">]</span></span></span></em>, <em class="sig-param"><span class="n"><span class="pre">num_mask_units</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">]</span></span></span></em>, <em class="sig-param"><span class="n"><span class="pre">architecture</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">Dict</span><span class="p"><span class="pre">]</span></span></span></em>, <em class="sig-param"><span class="n"><span class="pre">emb_dim</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">64</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">spatial_dims</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">3</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">patch_size</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">]</span></span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">4</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">context_pixels</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">]</span></span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">0</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">input_channels</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">1</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">save_layers</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">bool</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">False</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/encoder.html#HieraEncoder"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.encoder.HieraEncoder" title="Permalink to this definition">#</a></dt>
+<dd><p>Bases: <code class="xref py py-class docutils literal notranslate"><span class="pre">Module</span></code></p>
+<dl class="field-list">
+<dt class="field-odd">Parameters<span class="colon">:</span></dt>
+<dd class="field-odd"><ul>
+<li><p><strong>num_patches</strong> (<em>int, List[int]</em>) – Number of patches in each dimension. If a single int is provided, the number of patches in each dimension will be the same.</p></li>
+<li><p><strong>num_mask_units</strong> (<em>int, List[int]</em>) – Number of mask units in each dimension. If a single int is provided, the number of mask units in each dimension will be the same.</p></li>
+<li><p><strong>architecture</strong> (<em>List[Dict]</em>) – List of dictionaries specifying the architecture of the transformer. Each dictionary should have the following keys:
+- repeat: int</p>
+<blockquote>
+<div><p>Number of times to repeat the block</p>
+</div></blockquote>
+<ul class="simple">
+<li><dl class="simple">
+<dt>num_heads: int</dt><dd><p>Number of heads in the multihead attention</p>
+</dd>
+</dl>
+</li>
+<li><dl class="simple">
+<dt>q_stride: int, List[int]</dt><dd><p>Stride for the query in each spatial dimension</p>
+</dd>
+</dl>
+</li>
+<li><dl class="simple">
+<dt>self_attention: bool</dt><dd><p>Whether to use self attention or mask unit attention</p>
+</dd>
+</dl>
+</li>
+</ul>
+<p>On the last repeat of each non-self-attention block, the embedding dimension is doubled and spatial pooling with <cite>q_stride</cite> is performed within each mask unit. For example, a block with a embed_dim=4, q_stride=2, and repeat=2, the first repeat just does mask unit attention, while the second will produce an 8-dimensional output that has been spatially pooled.</p>
+</li>
+<li><p><strong>emb_dim</strong> (<em>int</em>) – Dimension of embedding</p></li>
+<li><p><strong>spatial_dims</strong> (<em>int</em>) – Number of spatial dimensions</p></li>
+<li><p><strong>patch_size</strong> (<em>List[int]</em>) – Size of each patch</p></li>
+<li><p><strong>context_pixels</strong> (<em>List[int]</em>) – Number of extra pixels around each patch to include in convolutional embedding to encoder dimension.</p></li>
+<li><p><strong>input_channels</strong> (<em>int</em>) – Number of input channels</p></li>
+<li><p><strong>save_layers</strong> (<em>bool</em>) – Whether to save the intermediate layer outputs</p></li>
+</ul>
+</dd>
+</dl>
+<dl class="py method">
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.encoder.HieraEncoder.forward">
+<span class="sig-name descname"><span class="pre">forward</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">img</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">mask_ratio</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/encoder.html#HieraEncoder.forward"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.encoder.HieraEncoder.forward" title="Permalink to this definition">#</a></dt>
+<dd></dd></dl>
+
+</dd></dl>
+
+<dl class="py class">
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.encoder.JEPAEncoder">
+<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">cyto_dl.nn.vits.encoder.</span></span><span class="sig-name descname"><span class="pre">JEPAEncoder</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">num_patches</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">]</span></span></span></em>, <em class="sig-param"><span class="n"><span class="pre">spatial_dims</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">3</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">patch_size</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">]</span></span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">4</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">emb_dim</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">192</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">num_layer</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">12</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">num_head</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">3</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">context_pixels</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">]</span></span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">0</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">input_channels</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">1</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">learnable_pos_embedding</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">bool</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">True</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/encoder.html#JEPAEncoder"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.encoder.JEPAEncoder" title="Permalink to this definition">#</a></dt>
+<dd><p>Bases: <code class="xref py py-class docutils literal notranslate"><span class="pre">Module</span></code></p>
+<dl class="field-list simple">
+<dt class="field-odd">Parameters<span class="colon">:</span></dt>
+<dd class="field-odd"><ul class="simple">
+<li><p><strong>num_patches</strong> (<em>List[int], int</em>) – Number of patches in each dimension. If a single int is provided, the number of patches in each dimension will be the same.</p></li>
+<li><p><strong>spatial_dims</strong> (<em>int</em>) – Number of spatial dimensions</p></li>
+<li><p><strong>patch_size</strong> (<em>List[int]</em>) – Size of each patch</p></li>
+<li><p><strong>emb_dim</strong> (<em>int</em>) – Dimension of embedding</p></li>
+<li><p><strong>num_layer</strong> (<em>int</em>) – Number of transformer layers</p></li>
+<li><p><strong>num_head</strong> (<em>int</em>) – Number of heads in transformer</p></li>
+<li><p><strong>context_pixels</strong> (<em>List[int], int</em>) – Number of extra pixels around each patch to include in convolutional embedding to encoder dimension. If a single int is provided, the number of context pixels in each dimension will be the same.</p></li>
+<li><p><strong>input_channels</strong> (<em>int</em>) – Number of input channels</p></li>
+<li><p><strong>learnable_pos_embedding</strong> (<em>bool</em>) – If True, learnable positional embeddings are used. If False, fixed sin/cos positional embeddings. Empirically, fixed positional embeddings work better for brightfield images.</p></li>
+</ul>
+</dd>
+</dl>
+<dl class="py method">
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.encoder.JEPAEncoder.forward">
+<span class="sig-name descname"><span class="pre">forward</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">img</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">patchify</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">True</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/encoder.html#JEPAEncoder.forward"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.encoder.JEPAEncoder.forward" title="Permalink to this definition">#</a></dt>
+<dd></dd></dl>
+
+</dd></dl>
+
+<dl class="py class">
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.encoder.MAE_Encoder">
+<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">cyto_dl.nn.vits.encoder.</span></span><span class="sig-name descname"><span class="pre">MAE_Encoder</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">num_patches</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">]</span></span></span></em>, <em class="sig-param"><span class="n"><span class="pre">spatial_dims</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">3</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">patch_size</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">]</span></span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">4</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">emb_dim</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">192</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">num_layer</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">12</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">num_head</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">3</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">context_pixels</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">]</span></span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">0</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">input_channels</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">1</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">n_intermediate_weights</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">-1</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/encoder.html#MAE_Encoder"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.encoder.MAE_Encoder" title="Permalink to this definition">#</a></dt>
+<dd><p>Bases: <code class="xref py py-class docutils literal notranslate"><span class="pre">Module</span></code></p>
+<dl class="field-list simple">
+<dt class="field-odd">Parameters<span class="colon">:</span></dt>
+<dd class="field-odd"><ul class="simple">
+<li><p><strong>num_patches</strong> (<em>List[int], int</em>) – Number of patches in each dimension. If a single int is provided, the number of patches in each dimension will be the same.</p></li>
+<li><p><strong>spatial_dims</strong> (<em>int</em>) – Number of spatial dimensions</p></li>
+<li><p><strong>patch_size</strong> (<em>List[int]</em>) – Size of each patch</p></li>
+<li><p><strong>emb_dim</strong> (<em>int</em>) – Dimension of embedding</p></li>
+<li><p><strong>num_layer</strong> (<em>int</em>) – Number of transformer layers</p></li>
+<li><p><strong>num_head</strong> (<em>int</em>) – Number of heads in transformer</p></li>
+<li><p><strong>context_pixels</strong> (<em>List[int], int</em>) – Number of extra pixels around each patch to include in convolutional embedding to encoder dimension. If a single int is provided, the number of context pixels in each dimension will be the same.</p></li>
+<li><p><strong>input_channels</strong> (<em>int</em>) – Number of input channels</p></li>
+<li><p><strong>n_intermediate_weights</strong> (<em>bool</em>) – Whether to use intermediate weights for weighted sum of intermediate layers</p></li>
+</ul>
+</dd>
+</dl>
+<dl class="py method">
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.encoder.MAE_Encoder.forward">
+<span class="sig-name descname"><span class="pre">forward</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">img</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">mask_ratio</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">0.75</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/encoder.html#MAE_Encoder.forward"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.encoder.MAE_Encoder.forward" title="Permalink to this definition">#</a></dt>
+<dd></dd></dl>
+
+<dl class="py method">
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.encoder.MAE_Encoder.init_weight">
+<span class="sig-name descname"><span class="pre">init_weight</span></span><span class="sig-paren">(</span><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/encoder.html#MAE_Encoder.init_weight"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.encoder.MAE_Encoder.init_weight" title="Permalink to this definition">#</a></dt>
+<dd></dd></dl>
+
+</dd></dl>
+
+<dl class="py class">
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.encoder.SpatialMerger">
+<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">cyto_dl.nn.vits.encoder.</span></span><span class="sig-name descname"><span class="pre">SpatialMerger</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">downsample_factor</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">]</span></span></span></em>, <em class="sig-param"><span class="n"><span class="pre">in_dim</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">out_dim</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">spatial_dims</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">3</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/encoder.html#SpatialMerger"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.encoder.SpatialMerger" title="Permalink to this definition">#</a></dt>
+<dd><p>Bases: <code class="xref py py-class docutils literal notranslate"><span class="pre">Module</span></code></p>
+<p>Class for converting multi-resolution Hiera features to the same (lowest) spatial resolution
+via convolution.</p>
+<dl class="py method">
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.encoder.SpatialMerger.forward">
+<span class="sig-name descname"><span class="pre">forward</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">x</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/encoder.html#SpatialMerger.forward"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.encoder.SpatialMerger.forward" title="Permalink to this definition">#</a></dt>
+<dd></dd></dl>
+
+</dd></dl>
+
+</section>
+
+        </article>
+      </div>
+      <footer>
+        
+        <div class="related-pages">
+          <a class="next-page" href="cyto_dl.nn.vits.mae.html">
+              <div class="page-info">
+                <div class="context">
+                  <span>Next</span>
+                </div>
+                <div class="title">cyto_dl.nn.vits.mae module</div>
+              </div>
+              <svg class="furo-related-icon"><use href="#svg-arrow-right"></use></svg>
+            </a>
+          <a class="prev-page" href="cyto_dl.nn.vits.decoder.html">
+              <svg class="furo-related-icon"><use href="#svg-arrow-right"></use></svg>
+              <div class="page-info">
+                <div class="context">
+                  <span>Previous</span>
+                </div>
+                
+                <div class="title">cyto_dl.nn.vits.decoder module</div>
+                
+              </div>
+            </a>
+        </div>
+        <div class="bottom-of-page">
+          <div class="left-details">
+            <div class="copyright">
+                Copyright &#169; 2023
+            </div>
+            Made with <a href="https://www.sphinx-doc.org/">Sphinx</a> and <a class="muted-link" href="https://pradyunsg.me">@pradyunsg</a>'s
+            
+            <a href="https://github.com/pradyunsg/furo">Furo</a>
+            
+          </div>
+          <div class="right-details">
+            
+          </div>
+        </div>
+        
+      </footer>
+    </div>
+    <aside class="toc-drawer">
+      
+      
+      <div class="toc-sticky toc-scroll">
+        <div class="toc-title-container">
+          <span class="toc-title">
+            On this page
+          </span>
+        </div>
+        <div class="toc-tree-container">
+          <div class="toc-tree">
+            <ul>
+<li><a class="reference internal" href="#">cyto_dl.nn.vits.encoder module</a><ul>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.encoder.HieraEncoder"><code class="docutils literal notranslate"><span class="pre">HieraEncoder</span></code></a><ul>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.encoder.HieraEncoder.forward"><code class="docutils literal notranslate"><span class="pre">HieraEncoder.forward()</span></code></a></li>
+</ul>
+</li>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.encoder.JEPAEncoder"><code class="docutils literal notranslate"><span class="pre">JEPAEncoder</span></code></a><ul>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.encoder.JEPAEncoder.forward"><code class="docutils literal notranslate"><span class="pre">JEPAEncoder.forward()</span></code></a></li>
+</ul>
+</li>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.encoder.MAE_Encoder"><code class="docutils literal notranslate"><span class="pre">MAE_Encoder</span></code></a><ul>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.encoder.MAE_Encoder.forward"><code class="docutils literal notranslate"><span class="pre">MAE_Encoder.forward()</span></code></a></li>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.encoder.MAE_Encoder.init_weight"><code class="docutils literal notranslate"><span class="pre">MAE_Encoder.init_weight()</span></code></a></li>
+</ul>
+</li>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.encoder.SpatialMerger"><code class="docutils literal notranslate"><span class="pre">SpatialMerger</span></code></a><ul>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.encoder.SpatialMerger.forward"><code class="docutils literal notranslate"><span class="pre">SpatialMerger.forward()</span></code></a></li>
+</ul>
+</li>
+</ul>
+</li>
+</ul>
+
+          </div>
+        </div>
+      </div>
+      
+      
+    </aside>
+  </div>
+</div><script data-url_root="./" id="documentation_options" src="_static/documentation_options.js"></script>
+    <script src="_static/doctools.js"></script>
+    <script src="_static/sphinx_highlight.js"></script>
+    <script src="_static/scripts/furo.js"></script>
+    </body>
+</html>
\ No newline at end of file
diff --git a/cyto_dl.nn.vits.html b/cyto_dl.nn.vits.html
index c523e1431..b8f6370f9 100644
--- a/cyto_dl.nn.vits.html
+++ b/cyto_dl.nn.vits.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 current has-children current-page"><a class="current reference internal" href="#">cyto_dl.nn.vits package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
@@ -441,6 +448,13 @@ <h2>Subpackages<a class="headerlink" href="#subpackages" title="Permalink to thi
 <div class="toctree-wrapper compound">
 <ul>
 <li class="toctree-l1"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><ul>
+<li class="toctree-l2"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html#subpackages">Subpackages</a><ul>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><ul>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html#submodules">Submodules</a></li>
+</ul>
+</li>
+</ul>
+</li>
 <li class="toctree-l2"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html#submodules">Submodules</a><ul>
 <li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a><ul>
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html#cyto_dl.nn.vits.blocks.cross_attention.CrossAttention"><code class="docutils literal notranslate"><span class="pre">CrossAttention</span></code></a></li>
@@ -453,9 +467,9 @@ <h2>Subpackages<a class="headerlink" href="#subpackages" title="Permalink to thi
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html#cyto_dl.nn.vits.blocks.intermediate_weigher.IntermediateWeigher"><code class="docutils literal notranslate"><span class="pre">IntermediateWeigher</span></code></a></li>
 </ul>
 </li>
-<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a><ul>
-<li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html#cyto_dl.nn.vits.blocks.patchify.Patchify"><code class="docutils literal notranslate"><span class="pre">Patchify</span></code></a></li>
-<li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html#cyto_dl.nn.vits.blocks.patchify.random_indexes"><code class="docutils literal notranslate"><span class="pre">random_indexes()</span></code></a></li>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a><ul>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html#cyto_dl.nn.vits.blocks.masked_unit_attention.HieraBlock"><code class="docutils literal notranslate"><span class="pre">HieraBlock</span></code></a></li>
+<li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html#cyto_dl.nn.vits.blocks.masked_unit_attention.MaskUnitAttention"><code class="docutils literal notranslate"><span class="pre">MaskUnitAttention</span></code></a></li>
 </ul>
 </li>
 </ul>
@@ -469,44 +483,70 @@ <h2>Subpackages<a class="headerlink" href="#subpackages" title="Permalink to thi
 <h2>Submodules<a class="headerlink" href="#submodules" title="Permalink to this heading">#</a></h2>
 <div class="toctree-wrapper compound">
 <ul>
-<li class="toctree-l1"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a><ul>
-<li class="toctree-l2"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html#cyto_dl.nn.vits.cross_mae.CrossMAE_Decoder"><code class="docutils literal notranslate"><span class="pre">CrossMAE_Decoder</span></code></a><ul>
-<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html#cyto_dl.nn.vits.cross_mae.CrossMAE_Decoder.forward"><code class="docutils literal notranslate"><span class="pre">CrossMAE_Decoder.forward()</span></code></a></li>
-<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html#cyto_dl.nn.vits.cross_mae.CrossMAE_Decoder.init_weight"><code class="docutils literal notranslate"><span class="pre">CrossMAE_Decoder.init_weight()</span></code></a></li>
+<li class="toctree-l1"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a><ul>
+<li class="toctree-l2"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html#cyto_dl.nn.vits.decoder.CrossMAE_Decoder"><code class="docutils literal notranslate"><span class="pre">CrossMAE_Decoder</span></code></a><ul>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html#cyto_dl.nn.vits.decoder.CrossMAE_Decoder.forward"><code class="docutils literal notranslate"><span class="pre">CrossMAE_Decoder.forward()</span></code></a></li>
+</ul>
+</li>
+<li class="toctree-l2"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html#cyto_dl.nn.vits.decoder.MAE_Decoder"><code class="docutils literal notranslate"><span class="pre">MAE_Decoder</span></code></a><ul>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html#cyto_dl.nn.vits.decoder.MAE_Decoder.add_mask_tokens"><code class="docutils literal notranslate"><span class="pre">MAE_Decoder.add_mask_tokens()</span></code></a></li>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html#cyto_dl.nn.vits.decoder.MAE_Decoder.adjust_indices_for_cls"><code class="docutils literal notranslate"><span class="pre">MAE_Decoder.adjust_indices_for_cls()</span></code></a></li>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html#cyto_dl.nn.vits.decoder.MAE_Decoder.forward"><code class="docutils literal notranslate"><span class="pre">MAE_Decoder.forward()</span></code></a></li>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html#cyto_dl.nn.vits.decoder.MAE_Decoder.init_weight"><code class="docutils literal notranslate"><span class="pre">MAE_Decoder.init_weight()</span></code></a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l1"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a><ul>
-<li class="toctree-l2"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html#cyto_dl.nn.vits.jepa.IWMPredictor"><code class="docutils literal notranslate"><span class="pre">IWMPredictor</span></code></a><ul>
-<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html#cyto_dl.nn.vits.jepa.IWMPredictor.forward"><code class="docutils literal notranslate"><span class="pre">IWMPredictor.forward()</span></code></a></li>
+<li class="toctree-l1"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a><ul>
+<li class="toctree-l2"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html#cyto_dl.nn.vits.encoder.HieraEncoder"><code class="docutils literal notranslate"><span class="pre">HieraEncoder</span></code></a><ul>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html#cyto_dl.nn.vits.encoder.HieraEncoder.forward"><code class="docutils literal notranslate"><span class="pre">HieraEncoder.forward()</span></code></a></li>
 </ul>
 </li>
-<li class="toctree-l2"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html#cyto_dl.nn.vits.jepa.JEPAEncoder"><code class="docutils literal notranslate"><span class="pre">JEPAEncoder</span></code></a><ul>
-<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html#cyto_dl.nn.vits.jepa.JEPAEncoder.forward"><code class="docutils literal notranslate"><span class="pre">JEPAEncoder.forward()</span></code></a></li>
+<li class="toctree-l2"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html#cyto_dl.nn.vits.encoder.JEPAEncoder"><code class="docutils literal notranslate"><span class="pre">JEPAEncoder</span></code></a><ul>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html#cyto_dl.nn.vits.encoder.JEPAEncoder.forward"><code class="docutils literal notranslate"><span class="pre">JEPAEncoder.forward()</span></code></a></li>
 </ul>
 </li>
-<li class="toctree-l2"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html#cyto_dl.nn.vits.jepa.JEPAPredictor"><code class="docutils literal notranslate"><span class="pre">JEPAPredictor</span></code></a><ul>
-<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html#cyto_dl.nn.vits.jepa.JEPAPredictor.forward"><code class="docutils literal notranslate"><span class="pre">JEPAPredictor.forward()</span></code></a></li>
-<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html#cyto_dl.nn.vits.jepa.JEPAPredictor.init_weight"><code class="docutils literal notranslate"><span class="pre">JEPAPredictor.init_weight()</span></code></a></li>
-<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html#cyto_dl.nn.vits.jepa.JEPAPredictor.predict_target_features"><code class="docutils literal notranslate"><span class="pre">JEPAPredictor.predict_target_features()</span></code></a></li>
+<li class="toctree-l2"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html#cyto_dl.nn.vits.encoder.MAE_Encoder"><code class="docutils literal notranslate"><span class="pre">MAE_Encoder</span></code></a><ul>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html#cyto_dl.nn.vits.encoder.MAE_Encoder.forward"><code class="docutils literal notranslate"><span class="pre">MAE_Encoder.forward()</span></code></a></li>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html#cyto_dl.nn.vits.encoder.MAE_Encoder.init_weight"><code class="docutils literal notranslate"><span class="pre">MAE_Encoder.init_weight()</span></code></a></li>
+</ul>
+</li>
+<li class="toctree-l2"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html#cyto_dl.nn.vits.encoder.SpatialMerger"><code class="docutils literal notranslate"><span class="pre">SpatialMerger</span></code></a><ul>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html#cyto_dl.nn.vits.encoder.SpatialMerger.forward"><code class="docutils literal notranslate"><span class="pre">SpatialMerger.forward()</span></code></a></li>
 </ul>
 </li>
 </ul>
 </li>
 <li class="toctree-l1"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a><ul>
-<li class="toctree-l2"><a class="reference internal" href="cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.MAE_Decoder"><code class="docutils literal notranslate"><span class="pre">MAE_Decoder</span></code></a><ul>
-<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.MAE_Decoder.forward"><code class="docutils literal notranslate"><span class="pre">MAE_Decoder.forward()</span></code></a></li>
-<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.MAE_Decoder.init_weight"><code class="docutils literal notranslate"><span class="pre">MAE_Decoder.init_weight()</span></code></a></li>
+<li class="toctree-l2"><a class="reference internal" href="cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.HieraMAE"><code class="docutils literal notranslate"><span class="pre">HieraMAE</span></code></a><ul>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.HieraMAE.decoder"><code class="docutils literal notranslate"><span class="pre">HieraMAE.decoder</span></code></a></li>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.HieraMAE.encoder"><code class="docutils literal notranslate"><span class="pre">HieraMAE.encoder</span></code></a></li>
+</ul>
+</li>
+<li class="toctree-l2"><a class="reference internal" href="cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.MAE"><code class="docutils literal notranslate"><span class="pre">MAE</span></code></a><ul>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.MAE.decoder"><code class="docutils literal notranslate"><span class="pre">MAE.decoder</span></code></a></li>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.MAE.encoder"><code class="docutils literal notranslate"><span class="pre">MAE.encoder</span></code></a></li>
+</ul>
+</li>
+<li class="toctree-l2"><a class="reference internal" href="cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.MAE_Base"><code class="docutils literal notranslate"><span class="pre">MAE_Base</span></code></a><ul>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.MAE_Base.decoder"><code class="docutils literal notranslate"><span class="pre">MAE_Base.decoder</span></code></a></li>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.MAE_Base.encoder"><code class="docutils literal notranslate"><span class="pre">MAE_Base.encoder</span></code></a></li>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.MAE_Base.forward"><code class="docutils literal notranslate"><span class="pre">MAE_Base.forward()</span></code></a></li>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.MAE_Base.init_decoder"><code class="docutils literal notranslate"><span class="pre">MAE_Base.init_decoder()</span></code></a></li>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.MAE_Base.init_encoder"><code class="docutils literal notranslate"><span class="pre">MAE_Base.init_encoder()</span></code></a></li>
+</ul>
+</li>
 </ul>
 </li>
-<li class="toctree-l2"><a class="reference internal" href="cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.MAE_Encoder"><code class="docutils literal notranslate"><span class="pre">MAE_Encoder</span></code></a><ul>
-<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.MAE_Encoder.forward"><code class="docutils literal notranslate"><span class="pre">MAE_Encoder.forward()</span></code></a></li>
-<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.MAE_Encoder.init_weight"><code class="docutils literal notranslate"><span class="pre">MAE_Encoder.init_weight()</span></code></a></li>
+<li class="toctree-l1"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a><ul>
+<li class="toctree-l2"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html#cyto_dl.nn.vits.predictor.IWMPredictor"><code class="docutils literal notranslate"><span class="pre">IWMPredictor</span></code></a><ul>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html#cyto_dl.nn.vits.predictor.IWMPredictor.forward"><code class="docutils literal notranslate"><span class="pre">IWMPredictor.forward()</span></code></a></li>
 </ul>
 </li>
-<li class="toctree-l2"><a class="reference internal" href="cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.MAE_ViT"><code class="docutils literal notranslate"><span class="pre">MAE_ViT</span></code></a><ul>
-<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.MAE_ViT.forward"><code class="docutils literal notranslate"><span class="pre">MAE_ViT.forward()</span></code></a></li>
+<li class="toctree-l2"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html#cyto_dl.nn.vits.predictor.JEPAPredictor"><code class="docutils literal notranslate"><span class="pre">JEPAPredictor</span></code></a><ul>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html#cyto_dl.nn.vits.predictor.JEPAPredictor.forward"><code class="docutils literal notranslate"><span class="pre">JEPAPredictor.forward()</span></code></a></li>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html#cyto_dl.nn.vits.predictor.JEPAPredictor.init_weight"><code class="docutils literal notranslate"><span class="pre">JEPAPredictor.init_weight()</span></code></a></li>
+<li class="toctree-l3"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html#cyto_dl.nn.vits.predictor.JEPAPredictor.predict_target_features"><code class="docutils literal notranslate"><span class="pre">JEPAPredictor.predict_target_features()</span></code></a></li>
 </ul>
 </li>
 </ul>
@@ -528,6 +568,8 @@ <h2>Submodules<a class="headerlink" href="#submodules" title="Permalink to this
 </li>
 <li class="toctree-l1"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a><ul>
 <li class="toctree-l2"><a class="reference internal" href="cyto_dl.nn.vits.utils.html#cyto_dl.nn.vits.utils.get_positional_embedding"><code class="docutils literal notranslate"><span class="pre">get_positional_embedding()</span></code></a></li>
+<li class="toctree-l2"><a class="reference internal" href="cyto_dl.nn.vits.utils.html#cyto_dl.nn.vits.utils.match_tuple_dimensions"><code class="docutils literal notranslate"><span class="pre">match_tuple_dimensions()</span></code></a></li>
+<li class="toctree-l2"><a class="reference internal" href="cyto_dl.nn.vits.utils.html#cyto_dl.nn.vits.utils.random_indexes"><code class="docutils literal notranslate"><span class="pre">random_indexes()</span></code></a></li>
 <li class="toctree-l2"><a class="reference internal" href="cyto_dl.nn.vits.utils.html#cyto_dl.nn.vits.utils.take_indexes"><code class="docutils literal notranslate"><span class="pre">take_indexes()</span></code></a></li>
 </ul>
 </li>
diff --git a/cyto_dl.nn.vits.mae.html b/cyto_dl.nn.vits.mae.html
index 277c46691..86096565f 100644
--- a/cyto_dl.nn.vits.mae.html
+++ b/cyto_dl.nn.vits.mae.html
@@ -3,7 +3,7 @@
   <head><meta charset="utf-8"/>
     <meta name="viewport" content="width=device-width,initial-scale=1"/>
     <meta name="color-scheme" content="light dark"><meta name="generator" content="Docutils 0.19: https://docutils.sourceforge.io/" />
-<link rel="index" title="Index" href="genindex.html" /><link rel="search" title="Search" href="search.html" /><link rel="next" title="cyto_dl.nn.vits.seg module" href="cyto_dl.nn.vits.seg.html" /><link rel="prev" title="cyto_dl.nn.vits.jepa module" href="cyto_dl.nn.vits.jepa.html" />
+<link rel="index" title="Index" href="genindex.html" /><link rel="search" title="Search" href="search.html" /><link rel="next" title="cyto_dl.nn.vits.predictor module" href="cyto_dl.nn.vits.predictor.html" /><link rel="prev" title="cyto_dl.nn.vits.encoder module" href="cyto_dl.nn.vits.encoder.html" />
 
     <!-- Generated with Sphinx 6.1.3 and Furo 2023.03.23 -->
         <title>cyto_dl.nn.vits.mae module - cyto-dl 0.4.0 documentation</title>
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 current has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5 current current-page"><a class="current reference internal" href="#">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
@@ -437,75 +444,73 @@
           <section id="module-cyto_dl.nn.vits.mae">
 <span id="cyto-dl-nn-vits-mae-module"></span><h1>cyto_dl.nn.vits.mae module<a class="headerlink" href="#module-cyto_dl.nn.vits.mae" title="Permalink to this heading">#</a></h1>
 <dl class="py class">
-<dt class="sig sig-object py" id="cyto_dl.nn.vits.mae.MAE_Decoder">
-<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">cyto_dl.nn.vits.mae.</span></span><span class="sig-name descname"><span class="pre">MAE_Decoder</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">num_patches</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">]</span></span></span></em>, <em class="sig-param"><span class="n"><span class="pre">spatial_dims</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">3</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">base_patch_size</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">]</span></span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">[4,</span> <span class="pre">8,</span> <span class="pre">8]</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">enc_dim</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">768</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">emb_dim</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">192</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">num_layer</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">4</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">num_head</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">3</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">learnable_pos_embedding</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">bool</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">True</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/mae.html#MAE_Decoder"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.mae.MAE_Decoder" title="Permalink to this definition">#</a></dt>
-<dd><p>Bases: <code class="xref py py-class docutils literal notranslate"><span class="pre">Module</span></code></p>
-<dl class="field-list simple">
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.mae.HieraMAE">
+<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">cyto_dl.nn.vits.mae.</span></span><span class="sig-name descname"><span class="pre">HieraMAE</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">architecture</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">Dict</span><span class="p"><span class="pre">]</span></span></span></em>, <em class="sig-param"><span class="n"><span class="pre">spatial_dims</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">3</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">num_patches</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">]</span></span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">16</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">num_mask_units</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">]</span></span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">8</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">patch_size</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">]</span></span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">4</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">emb_dim</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">64</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">decoder_layer</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">4</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">decoder_head</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">8</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">decoder_dim</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">192</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">mask_ratio</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">0.75</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">use_crossmae</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">bool</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">False</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">context_pixels</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">]</span></span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">0</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">input_channels</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">1</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">features_only</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">bool</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">False</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/mae.html#HieraMAE"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.mae.HieraMAE" title="Permalink to this definition">#</a></dt>
+<dd><p>Bases: <a class="reference internal" href="#cyto_dl.nn.vits.mae.MAE_Base" title="cyto_dl.nn.vits.mae.MAE_Base"><code class="xref py py-class docutils literal notranslate"><span class="pre">MAE_Base</span></code></a></p>
+<dl class="field-list">
 <dt class="field-odd">Parameters<span class="colon">:</span></dt>
-<dd class="field-odd"><ul class="simple">
-<li><p><strong>num_patches</strong> (<em>List[int]</em>) – Number of patches in each dimension</p></li>
-<li><p><strong>base_patch_size</strong> (<em>Tuple[int]</em>) – Size of each patch</p></li>
-<li><p><strong>enc_dim</strong> (<em>int</em>) – Dimension of encoder</p></li>
-<li><p><strong>emb_dim</strong> (<em>int</em>) – Dimension of decoder</p></li>
-<li><p><strong>num_layer</strong> (<em>int</em>) – Number of transformer layers</p></li>
-<li><p><strong>num_head</strong> (<em>int</em>) – Number of heads in transformer</p></li>
-<li><p><strong>learnable_pos_embedding</strong> (<em>bool</em>) – If True, learnable positional embeddings are used. If False, fixed sin/cos positional embeddings. Empirically, fixed positional embeddings work better for brightfield images.</p></li>
-</ul>
+<dd class="field-odd"><ul>
+<li><p><strong>architecture</strong> (<em>List[Dict]</em>) – List of dictionaries specifying the architecture of the transformer. Each dictionary should have the following keys:
+- repeat: int</p>
+<blockquote>
+<div><p>Number of times to repeat the block</p>
+</div></blockquote>
+<ul class="simple">
+<li><dl class="simple">
+<dt>num_heads: int</dt><dd><p>Number of heads in the multihead attention</p>
 </dd>
 </dl>
-<dl class="py method">
-<dt class="sig sig-object py" id="cyto_dl.nn.vits.mae.MAE_Decoder.forward">
-<span class="sig-name descname"><span class="pre">forward</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">features</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">forward_indexes</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">backward_indexes</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/mae.html#MAE_Decoder.forward"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.mae.MAE_Decoder.forward" title="Permalink to this definition">#</a></dt>
-<dd></dd></dl>
-
-<dl class="py method">
-<dt class="sig sig-object py" id="cyto_dl.nn.vits.mae.MAE_Decoder.init_weight">
-<span class="sig-name descname"><span class="pre">init_weight</span></span><span class="sig-paren">(</span><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/mae.html#MAE_Decoder.init_weight"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.mae.MAE_Decoder.init_weight" title="Permalink to this definition">#</a></dt>
-<dd></dd></dl>
-
-</dd></dl>
-
-<dl class="py class">
-<dt class="sig sig-object py" id="cyto_dl.nn.vits.mae.MAE_Encoder">
-<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">cyto_dl.nn.vits.mae.</span></span><span class="sig-name descname"><span class="pre">MAE_Encoder</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">num_patches</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">]</span></span></span></em>, <em class="sig-param"><span class="n"><span class="pre">spatial_dims</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">3</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">base_patch_size</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">]</span></span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">(16,</span> <span class="pre">16,</span> <span class="pre">16)</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">emb_dim</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">192</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">num_layer</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">12</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">num_head</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">3</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">context_pixels</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">]</span></span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">[0,</span> <span class="pre">0,</span> <span class="pre">0]</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">input_channels</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">1</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">n_intermediate_weights</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">-1</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/mae.html#MAE_Encoder"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.mae.MAE_Encoder" title="Permalink to this definition">#</a></dt>
-<dd><p>Bases: <code class="xref py py-class docutils literal notranslate"><span class="pre">Module</span></code></p>
-<dl class="field-list simple">
-<dt class="field-odd">Parameters<span class="colon">:</span></dt>
-<dd class="field-odd"><ul class="simple">
-<li><p><strong>num_patches</strong> (<em>List[int]</em>) – Number of patches in each dimension</p></li>
+</li>
+<li><dl class="simple">
+<dt>q_stride: List[int]</dt><dd><p>Stride for the query in each spatial dimension</p>
+</dd>
+</dl>
+</li>
+<li><dl class="simple">
+<dt>self_attention: bool</dt><dd><p>Whether to use self attention or mask unit attention</p>
+</dd>
+</dl>
+</li>
+</ul>
+</li>
 <li><p><strong>spatial_dims</strong> (<em>int</em>) – Number of spatial dimensions</p></li>
-<li><p><strong>base_patch_size</strong> (<em>List[int]</em>) – Size of each patch</p></li>
+<li><p><strong>num_patches</strong> (<em>int, List[int]</em>) – Number of patches in each dimension (Z)YX order. If int, the same number of patches is used in each dimension.</p></li>
+<li><p><strong>num_mask_units</strong> (<em>int, List[int]</em>) – Number of mask units in each dimension (Z)YX order. If int, the same number of mask units is used in each dimension.</p></li>
+<li><p><strong>patch_size</strong> (<em>int, List[int]</em>) – Size of each patch (Z)YX order. If int, the same patch size is used in each dimension.</p></li>
 <li><p><strong>emb_dim</strong> (<em>int</em>) – Dimension of embedding</p></li>
-<li><p><strong>num_layer</strong> (<em>int</em>) – Number of transformer layers</p></li>
-<li><p><strong>num_head</strong> (<em>int</em>) – Number of heads in transformer</p></li>
+<li><p><strong>decoder_layer</strong> (<em>int</em>) – Number of layers in the decoder</p></li>
+<li><p><strong>decoder_head</strong> (<em>int</em>) – Number of heads in the decoder</p></li>
+<li><p><strong>decoder_dim</strong> (<em>int</em>) – Dimension of the decoder</p></li>
+<li><p><strong>mask_ratio</strong> (<em>float</em>) – Fraction of mask units to remove</p></li>
+<li><p><strong>use_crossmae</strong> (<em>bool</em>) – Use CrossMAE-style decoder instead of MAE decoder</p></li>
 <li><p><strong>context_pixels</strong> (<em>List[int]</em>) – Number of extra pixels around each patch to include in convolutional embedding to encoder dimension.</p></li>
 <li><p><strong>input_channels</strong> (<em>int</em>) – Number of input channels</p></li>
-<li><p><strong>weight_intermediates</strong> (<em>bool</em>) – Whether to output linear combination of intermediate layers as final output like CrossMAE</p></li>
+<li><p><strong>features_only</strong> (<em>bool</em>) – Only use encoder to extract features</p></li>
 </ul>
 </dd>
 </dl>
-<dl class="py method">
-<dt class="sig sig-object py" id="cyto_dl.nn.vits.mae.MAE_Encoder.forward">
-<span class="sig-name descname"><span class="pre">forward</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">img</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">mask_ratio</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">0.75</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/mae.html#MAE_Encoder.forward"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.mae.MAE_Encoder.forward" title="Permalink to this definition">#</a></dt>
+<dl class="py property">
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.mae.HieraMAE.decoder">
+<em class="property"><span class="pre">property</span><span class="w"> </span></em><span class="sig-name descname"><span class="pre">decoder</span></span><a class="headerlink" href="#cyto_dl.nn.vits.mae.HieraMAE.decoder" title="Permalink to this definition">#</a></dt>
 <dd></dd></dl>
 
-<dl class="py method">
-<dt class="sig sig-object py" id="cyto_dl.nn.vits.mae.MAE_Encoder.init_weight">
-<span class="sig-name descname"><span class="pre">init_weight</span></span><span class="sig-paren">(</span><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/mae.html#MAE_Encoder.init_weight"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.mae.MAE_Encoder.init_weight" title="Permalink to this definition">#</a></dt>
+<dl class="py property">
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.mae.HieraMAE.encoder">
+<em class="property"><span class="pre">property</span><span class="w"> </span></em><span class="sig-name descname"><span class="pre">encoder</span></span><a class="headerlink" href="#cyto_dl.nn.vits.mae.HieraMAE.encoder" title="Permalink to this definition">#</a></dt>
 <dd></dd></dl>
 
 </dd></dl>
 
 <dl class="py class">
-<dt class="sig sig-object py" id="cyto_dl.nn.vits.mae.MAE_ViT">
-<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">cyto_dl.nn.vits.mae.</span></span><span class="sig-name descname"><span class="pre">MAE_ViT</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">spatial_dims</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">3</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">num_patches</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">]</span></span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">[2,</span> <span class="pre">32,</span> <span class="pre">32]</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">base_patch_size</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">]</span></span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">[16,</span> <span class="pre">16,</span> <span class="pre">16]</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">emb_dim</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">768</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">encoder_layer</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">12</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">encoder_head</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">8</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">decoder_layer</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">4</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">decoder_head</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">8</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">decoder_dim</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">192</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">mask_ratio</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">0.75</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">use_crossmae</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">bool</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">False</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">context_pixels</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">]</span></span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">[0,</span> <span class="pre">0,</span> <span class="pre">0]</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">input_channels</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">1</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">features_only</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">bool</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">False</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">learnable_pos_embedding</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">bool</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">True</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/mae.html#MAE_ViT"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.mae.MAE_ViT" title="Permalink to this definition">#</a></dt>
-<dd><p>Bases: <code class="xref py py-class docutils literal notranslate"><span class="pre">Module</span></code></p>
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.mae.MAE">
+<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">cyto_dl.nn.vits.mae.</span></span><span class="sig-name descname"><span class="pre">MAE</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">spatial_dims</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">3</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">num_patches</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">]</span></span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">16</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">patch_size</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">]</span></span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">4</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">emb_dim</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">768</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">encoder_layer</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">12</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">encoder_head</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">8</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">decoder_layer</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">4</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">decoder_head</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">8</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">decoder_dim</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">192</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">mask_ratio</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">0.75</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">use_crossmae</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">bool</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">False</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">context_pixels</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">]</span></span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">0</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">input_channels</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">1</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">features_only</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">bool</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">False</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">learnable_pos_embedding</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">bool</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">True</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/mae.html#MAE"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.mae.MAE" title="Permalink to this definition">#</a></dt>
+<dd><p>Bases: <a class="reference internal" href="#cyto_dl.nn.vits.mae.MAE_Base" title="cyto_dl.nn.vits.mae.MAE_Base"><code class="xref py py-class docutils literal notranslate"><span class="pre">MAE_Base</span></code></a></p>
 <dl class="field-list simple">
 <dt class="field-odd">Parameters<span class="colon">:</span></dt>
 <dd class="field-odd"><ul class="simple">
 <li><p><strong>spatial_dims</strong> (<em>int</em>) – Number of spatial dimensions</p></li>
 <li><p><strong>num_patches</strong> (<em>List[int]</em>) – Number of patches in each dimension (ZYX order)</p></li>
-<li><p><strong>base_patch_size</strong> (<em>List[int]</em>) – Size of each patch (ZYX order)</p></li>
+<li><p><strong>patch_size</strong> (<em>List[int]</em>) – Size of each patch (ZYX order)</p></li>
 <li><p><strong>emb_dim</strong> (<em>int</em>) – Dimension of encoder embedding</p></li>
 <li><p><strong>encoder_layer</strong> (<em>int</em>) – Number of encoder transformer layers</p></li>
 <li><p><strong>encoder_head</strong> (<em>int</em>) – Number of encoder heads</p></li>
@@ -520,9 +525,45 @@
 </ul>
 </dd>
 </dl>
+<dl class="py property">
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.mae.MAE.decoder">
+<em class="property"><span class="pre">property</span><span class="w"> </span></em><span class="sig-name descname"><span class="pre">decoder</span></span><a class="headerlink" href="#cyto_dl.nn.vits.mae.MAE.decoder" title="Permalink to this definition">#</a></dt>
+<dd></dd></dl>
+
+<dl class="py property">
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.mae.MAE.encoder">
+<em class="property"><span class="pre">property</span><span class="w"> </span></em><span class="sig-name descname"><span class="pre">encoder</span></span><a class="headerlink" href="#cyto_dl.nn.vits.mae.MAE.encoder" title="Permalink to this definition">#</a></dt>
+<dd></dd></dl>
+
+</dd></dl>
+
+<dl class="py class">
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.mae.MAE_Base">
+<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">cyto_dl.nn.vits.mae.</span></span><span class="sig-name descname"><span class="pre">MAE_Base</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">spatial_dims</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">num_patches</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">patch_size</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">mask_ratio</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">features_only</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">context_pixels</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/mae.html#MAE_Base"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.mae.MAE_Base" title="Permalink to this definition">#</a></dt>
+<dd><p>Bases: <code class="xref py py-class docutils literal notranslate"><span class="pre">Module</span></code>, <code class="xref py py-class docutils literal notranslate"><span class="pre">ABC</span></code></p>
+<dl class="py property">
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.mae.MAE_Base.decoder">
+<em class="property"><span class="pre">abstract</span><span class="w"> </span><span class="pre">property</span><span class="w"> </span></em><span class="sig-name descname"><span class="pre">decoder</span></span><a class="headerlink" href="#cyto_dl.nn.vits.mae.MAE_Base.decoder" title="Permalink to this definition">#</a></dt>
+<dd></dd></dl>
+
+<dl class="py property">
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.mae.MAE_Base.encoder">
+<em class="property"><span class="pre">abstract</span><span class="w"> </span><span class="pre">property</span><span class="w"> </span></em><span class="sig-name descname"><span class="pre">encoder</span></span><a class="headerlink" href="#cyto_dl.nn.vits.mae.MAE_Base.encoder" title="Permalink to this definition">#</a></dt>
+<dd></dd></dl>
+
+<dl class="py method">
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.mae.MAE_Base.forward">
+<span class="sig-name descname"><span class="pre">forward</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">img</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/mae.html#MAE_Base.forward"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.mae.MAE_Base.forward" title="Permalink to this definition">#</a></dt>
+<dd></dd></dl>
+
+<dl class="py method">
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.mae.MAE_Base.init_decoder">
+<span class="sig-name descname"><span class="pre">init_decoder</span></span><span class="sig-paren">(</span><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/mae.html#MAE_Base.init_decoder"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.mae.MAE_Base.init_decoder" title="Permalink to this definition">#</a></dt>
+<dd></dd></dl>
+
 <dl class="py method">
-<dt class="sig sig-object py" id="cyto_dl.nn.vits.mae.MAE_ViT.forward">
-<span class="sig-name descname"><span class="pre">forward</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">img</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/mae.html#MAE_ViT.forward"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.mae.MAE_ViT.forward" title="Permalink to this definition">#</a></dt>
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.mae.MAE_Base.init_encoder">
+<span class="sig-name descname"><span class="pre">init_encoder</span></span><span class="sig-paren">(</span><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/mae.html#MAE_Base.init_encoder"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.mae.MAE_Base.init_encoder" title="Permalink to this definition">#</a></dt>
 <dd></dd></dl>
 
 </dd></dl>
@@ -534,23 +575,23 @@
       <footer>
         
         <div class="related-pages">
-          <a class="next-page" href="cyto_dl.nn.vits.seg.html">
+          <a class="next-page" href="cyto_dl.nn.vits.predictor.html">
               <div class="page-info">
                 <div class="context">
                   <span>Next</span>
                 </div>
-                <div class="title">cyto_dl.nn.vits.seg module</div>
+                <div class="title">cyto_dl.nn.vits.predictor module</div>
               </div>
               <svg class="furo-related-icon"><use href="#svg-arrow-right"></use></svg>
             </a>
-          <a class="prev-page" href="cyto_dl.nn.vits.jepa.html">
+          <a class="prev-page" href="cyto_dl.nn.vits.encoder.html">
               <svg class="furo-related-icon"><use href="#svg-arrow-right"></use></svg>
               <div class="page-info">
                 <div class="context">
                   <span>Previous</span>
                 </div>
                 
-                <div class="title">cyto_dl.nn.vits.jepa module</div>
+                <div class="title">cyto_dl.nn.vits.encoder module</div>
                 
               </div>
             </a>
@@ -585,18 +626,22 @@
           <div class="toc-tree">
             <ul>
 <li><a class="reference internal" href="#">cyto_dl.nn.vits.mae module</a><ul>
-<li><a class="reference internal" href="#cyto_dl.nn.vits.mae.MAE_Decoder"><code class="docutils literal notranslate"><span class="pre">MAE_Decoder</span></code></a><ul>
-<li><a class="reference internal" href="#cyto_dl.nn.vits.mae.MAE_Decoder.forward"><code class="docutils literal notranslate"><span class="pre">MAE_Decoder.forward()</span></code></a></li>
-<li><a class="reference internal" href="#cyto_dl.nn.vits.mae.MAE_Decoder.init_weight"><code class="docutils literal notranslate"><span class="pre">MAE_Decoder.init_weight()</span></code></a></li>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.mae.HieraMAE"><code class="docutils literal notranslate"><span class="pre">HieraMAE</span></code></a><ul>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.mae.HieraMAE.decoder"><code class="docutils literal notranslate"><span class="pre">HieraMAE.decoder</span></code></a></li>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.mae.HieraMAE.encoder"><code class="docutils literal notranslate"><span class="pre">HieraMAE.encoder</span></code></a></li>
 </ul>
 </li>
-<li><a class="reference internal" href="#cyto_dl.nn.vits.mae.MAE_Encoder"><code class="docutils literal notranslate"><span class="pre">MAE_Encoder</span></code></a><ul>
-<li><a class="reference internal" href="#cyto_dl.nn.vits.mae.MAE_Encoder.forward"><code class="docutils literal notranslate"><span class="pre">MAE_Encoder.forward()</span></code></a></li>
-<li><a class="reference internal" href="#cyto_dl.nn.vits.mae.MAE_Encoder.init_weight"><code class="docutils literal notranslate"><span class="pre">MAE_Encoder.init_weight()</span></code></a></li>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.mae.MAE"><code class="docutils literal notranslate"><span class="pre">MAE</span></code></a><ul>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.mae.MAE.decoder"><code class="docutils literal notranslate"><span class="pre">MAE.decoder</span></code></a></li>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.mae.MAE.encoder"><code class="docutils literal notranslate"><span class="pre">MAE.encoder</span></code></a></li>
 </ul>
 </li>
-<li><a class="reference internal" href="#cyto_dl.nn.vits.mae.MAE_ViT"><code class="docutils literal notranslate"><span class="pre">MAE_ViT</span></code></a><ul>
-<li><a class="reference internal" href="#cyto_dl.nn.vits.mae.MAE_ViT.forward"><code class="docutils literal notranslate"><span class="pre">MAE_ViT.forward()</span></code></a></li>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.mae.MAE_Base"><code class="docutils literal notranslate"><span class="pre">MAE_Base</span></code></a><ul>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.mae.MAE_Base.decoder"><code class="docutils literal notranslate"><span class="pre">MAE_Base.decoder</span></code></a></li>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.mae.MAE_Base.encoder"><code class="docutils literal notranslate"><span class="pre">MAE_Base.encoder</span></code></a></li>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.mae.MAE_Base.forward"><code class="docutils literal notranslate"><span class="pre">MAE_Base.forward()</span></code></a></li>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.mae.MAE_Base.init_decoder"><code class="docutils literal notranslate"><span class="pre">MAE_Base.init_decoder()</span></code></a></li>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.mae.MAE_Base.init_encoder"><code class="docutils literal notranslate"><span class="pre">MAE_Base.init_encoder()</span></code></a></li>
 </ul>
 </li>
 </ul>
diff --git a/cyto_dl.nn.vits.jepa.html b/cyto_dl.nn.vits.predictor.html
similarity index 73%
rename from cyto_dl.nn.vits.jepa.html
rename to cyto_dl.nn.vits.predictor.html
index 9681cee01..b1a059502 100644
--- a/cyto_dl.nn.vits.jepa.html
+++ b/cyto_dl.nn.vits.predictor.html
@@ -3,10 +3,10 @@
   <head><meta charset="utf-8"/>
     <meta name="viewport" content="width=device-width,initial-scale=1"/>
     <meta name="color-scheme" content="light dark"><meta name="generator" content="Docutils 0.19: https://docutils.sourceforge.io/" />
-<link rel="index" title="Index" href="genindex.html" /><link rel="search" title="Search" href="search.html" /><link rel="next" title="cyto_dl.nn.vits.mae module" href="cyto_dl.nn.vits.mae.html" /><link rel="prev" title="cyto_dl.nn.vits.cross_mae module" href="cyto_dl.nn.vits.cross_mae.html" />
+<link rel="index" title="Index" href="genindex.html" /><link rel="search" title="Search" href="search.html" /><link rel="next" title="cyto_dl.nn.vits.seg module" href="cyto_dl.nn.vits.seg.html" /><link rel="prev" title="cyto_dl.nn.vits.mae module" href="cyto_dl.nn.vits.mae.html" />
 
     <!-- Generated with Sphinx 6.1.3 and Furo 2023.03.23 -->
-        <title>cyto_dl.nn.vits.jepa module - cyto-dl 0.4.0 documentation</title>
+        <title>cyto_dl.nn.vits.predictor module - cyto-dl 0.4.0 documentation</title>
       <link rel="stylesheet" type="text/css" href="_static/pygments.css" />
     <link rel="stylesheet" type="text/css" href="_static/styles/furo.css?digest=fad236701ea90a88636c2a8c73b44ae642ed2a53" />
     <link rel="stylesheet" type="text/css" href="_static/styles/furo-extensions.css?digest=30d1aed668e5c3a91c3e3bf6a60b675221979f0e" />
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 current has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5 current current-page"><a class="current reference internal" href="#">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5 current current-page"><a class="current reference internal" href="#">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
@@ -434,67 +441,45 @@
           </label>
         </div>
         <article role="main">
-          <section id="module-cyto_dl.nn.vits.jepa">
-<span id="cyto-dl-nn-vits-jepa-module"></span><h1>cyto_dl.nn.vits.jepa module<a class="headerlink" href="#module-cyto_dl.nn.vits.jepa" title="Permalink to this heading">#</a></h1>
+          <section id="module-cyto_dl.nn.vits.predictor">
+<span id="cyto-dl-nn-vits-predictor-module"></span><h1>cyto_dl.nn.vits.predictor module<a class="headerlink" href="#module-cyto_dl.nn.vits.predictor" title="Permalink to this heading">#</a></h1>
 <dl class="py class">
-<dt class="sig sig-object py" id="cyto_dl.nn.vits.jepa.IWMPredictor">
-<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">cyto_dl.nn.vits.jepa.</span></span><span class="sig-name descname"><span class="pre">IWMPredictor</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">domains</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">str</span><span class="p"><span class="pre">]</span></span></span></em>, <em class="sig-param"><span class="n"><span class="pre">num_patches</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">]</span></span></span></em>, <em class="sig-param"><span class="n"><span class="pre">input_dim</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">192</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">emb_dim</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">192</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">num_layer</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">12</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">num_head</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">3</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/jepa.html#IWMPredictor"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.jepa.IWMPredictor" title="Permalink to this definition">#</a></dt>
-<dd><p>Bases: <a class="reference internal" href="#cyto_dl.nn.vits.jepa.JEPAPredictor" title="cyto_dl.nn.vits.jepa.JEPAPredictor"><code class="xref py py-class docutils literal notranslate"><span class="pre">JEPAPredictor</span></code></a></p>
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.predictor.IWMPredictor">
+<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">cyto_dl.nn.vits.predictor.</span></span><span class="sig-name descname"><span class="pre">IWMPredictor</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">domains</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">str</span><span class="p"><span class="pre">]</span></span></span></em>, <em class="sig-param"><span class="n"><span class="pre">num_patches</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">]</span></span></span></em>, <em class="sig-param"><span class="n"><span class="pre">spatial_dims</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">3</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">input_dim</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">192</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">emb_dim</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">192</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">num_layer</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">12</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">num_head</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">3</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/predictor.html#IWMPredictor"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.predictor.IWMPredictor" title="Permalink to this definition">#</a></dt>
+<dd><p>Bases: <a class="reference internal" href="#cyto_dl.nn.vits.predictor.JEPAPredictor" title="cyto_dl.nn.vits.predictor.JEPAPredictor"><code class="xref py py-class docutils literal notranslate"><span class="pre">JEPAPredictor</span></code></a></p>
 <p>Specialized JEPA predictor that can conditionally predict between different domains (e.g.
 predict from brightfield to multiple fluorescent tags)</p>
 <dl class="field-list simple">
 <dt class="field-odd">Parameters<span class="colon">:</span></dt>
 <dd class="field-odd"><ul class="simple">
 <li><p><strong>domains</strong> (<em>List[str]</em>) – List of names of target domains</p></li>
-<li><p><strong>num_patches</strong> (<em>List[int]</em>) – Number of patches in each dimension</p></li>
-<li><p><strong>emb_dim</strong> (<em>int</em>) – Dimension of embedding</p></li>
-<li><p><strong>num_layer</strong> (<em>int</em>) – Number of transformer layers</p></li>
-<li><p><strong>num_head</strong> (<em>int</em>) – Number of heads in transformer</p></li>
-</ul>
-</dd>
-</dl>
-<dl class="py method">
-<dt class="sig sig-object py" id="cyto_dl.nn.vits.jepa.IWMPredictor.forward">
-<span class="sig-name descname"><span class="pre">forward</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">context_emb</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">target_masks</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">target_domain</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/jepa.html#IWMPredictor.forward"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.jepa.IWMPredictor.forward" title="Permalink to this definition">#</a></dt>
-<dd></dd></dl>
-
-</dd></dl>
-
-<dl class="py class">
-<dt class="sig sig-object py" id="cyto_dl.nn.vits.jepa.JEPAEncoder">
-<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">cyto_dl.nn.vits.jepa.</span></span><span class="sig-name descname"><span class="pre">JEPAEncoder</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">num_patches</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">]</span></span></span></em>, <em class="sig-param"><span class="n"><span class="pre">spatial_dims</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">3</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">patch_size</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">]</span></span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">(16,</span> <span class="pre">16,</span> <span class="pre">16)</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">emb_dim</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">192</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">num_layer</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">12</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">num_head</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">3</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">context_pixels</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">]</span></span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">[0,</span> <span class="pre">0,</span> <span class="pre">0]</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">input_channels</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">1</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">learnable_pos_embedding</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">bool</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">True</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/jepa.html#JEPAEncoder"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.jepa.JEPAEncoder" title="Permalink to this definition">#</a></dt>
-<dd><p>Bases: <code class="xref py py-class docutils literal notranslate"><span class="pre">Module</span></code></p>
-<dl class="field-list simple">
-<dt class="field-odd">Parameters<span class="colon">:</span></dt>
-<dd class="field-odd"><ul class="simple">
-<li><p><strong>num_patches</strong> (<em>List[int]</em>) – Number of patches in each dimension</p></li>
+<li><p><strong>num_patches</strong> (<em>List[int]</em>) – Number of patches in each dimension. If int, the same number of patches is used for all spatial dimensions</p></li>
+<li><p><strong>spatial_dims</strong> (<em>int</em>) – Number of spatial dimensions</p></li>
 <li><p><strong>spatial_dims</strong> (<em>int</em>) – Number of spatial dimensions</p></li>
-<li><p><strong>patch_size</strong> (<em>List[int]</em>) – Size of each patch</p></li>
 <li><p><strong>emb_dim</strong> (<em>int</em>) – Dimension of embedding</p></li>
 <li><p><strong>num_layer</strong> (<em>int</em>) – Number of transformer layers</p></li>
 <li><p><strong>num_head</strong> (<em>int</em>) – Number of heads in transformer</p></li>
-<li><p><strong>context_pixels</strong> (<em>List[int]</em>) – Number of extra pixels around each patch to include in convolutional embedding to encoder dimension.</p></li>
-<li><p><strong>input_channels</strong> (<em>int</em>) – Number of input channels</p></li>
-<li><p><strong>learnable_pos_embedding</strong> (<em>bool</em>) – If True, learnable positional embeddings are used. If False, fixed sin/cos positional embeddings. Empirically, fixed positional embeddings work better for brightfield images.</p></li>
 </ul>
 </dd>
 </dl>
 <dl class="py method">
-<dt class="sig sig-object py" id="cyto_dl.nn.vits.jepa.JEPAEncoder.forward">
-<span class="sig-name descname"><span class="pre">forward</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">img</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">patchify</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">True</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/jepa.html#JEPAEncoder.forward"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.jepa.JEPAEncoder.forward" title="Permalink to this definition">#</a></dt>
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.predictor.IWMPredictor.forward">
+<span class="sig-name descname"><span class="pre">forward</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">context_emb</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">target_masks</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">target_domain</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/predictor.html#IWMPredictor.forward"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.predictor.IWMPredictor.forward" title="Permalink to this definition">#</a></dt>
 <dd></dd></dl>
 
 </dd></dl>
 
 <dl class="py class">
-<dt class="sig sig-object py" id="cyto_dl.nn.vits.jepa.JEPAPredictor">
-<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">cyto_dl.nn.vits.jepa.</span></span><span class="sig-name descname"><span class="pre">JEPAPredictor</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">num_patches</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">]</span></span></span></em>, <em class="sig-param"><span class="n"><span class="pre">input_dim</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">192</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">emb_dim</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">192</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">num_layer</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">12</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">num_head</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">3</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">learnable_pos_embedding</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">bool</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">True</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/jepa.html#JEPAPredictor"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.jepa.JEPAPredictor" title="Permalink to this definition">#</a></dt>
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.predictor.JEPAPredictor">
+<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">cyto_dl.nn.vits.predictor.</span></span><span class="sig-name descname"><span class="pre">JEPAPredictor</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">num_patches</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">int</span><span class="p"><span class="pre">]</span></span></span></em>, <em class="sig-param"><span class="n"><span class="pre">spatial_dims</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">3</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">input_dim</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">192</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">emb_dim</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">192</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">num_layer</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">12</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">num_head</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">3</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">learnable_pos_embedding</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">bool</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">True</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/predictor.html#JEPAPredictor"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.predictor.JEPAPredictor" title="Permalink to this definition">#</a></dt>
 <dd><p>Bases: <code class="xref py py-class docutils literal notranslate"><span class="pre">Module</span></code></p>
 <p>Class for predicting target features from context embedding.</p>
 <dl class="field-list simple">
 <dt class="field-odd">Parameters<span class="colon">:</span></dt>
 <dd class="field-odd"><ul class="simple">
-<li><p><strong>num_patches</strong> (<em>List[int]</em>) – Number of patches in each dimension</p></li>
+<li><p><strong>num_patches</strong> (<em>List[int], int</em>) – Number of patches in each dimension. If int, the same number of patches is used for all spatial dimensions</p></li>
+<li><p><strong>spatial_dims</strong> (<em>int</em>) – Number of spatial dimensions</p></li>
+<li><p><strong>input_dim</strong> (<em>int</em>) – Dimension of input</p></li>
 <li><p><strong>emb_dim</strong> (<em>int</em>) – Dimension of embedding</p></li>
 <li><p><strong>num_layer</strong> (<em>int</em>) – Number of transformer layers</p></li>
 <li><p><strong>num_head</strong> (<em>int</em>) – Number of heads in transformer</p></li>
@@ -503,18 +488,18 @@
 </dd>
 </dl>
 <dl class="py method">
-<dt class="sig sig-object py" id="cyto_dl.nn.vits.jepa.JEPAPredictor.forward">
-<span class="sig-name descname"><span class="pre">forward</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">context_emb</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">target_masks</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/jepa.html#JEPAPredictor.forward"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.jepa.JEPAPredictor.forward" title="Permalink to this definition">#</a></dt>
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.predictor.JEPAPredictor.forward">
+<span class="sig-name descname"><span class="pre">forward</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">context_emb</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">target_masks</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/predictor.html#JEPAPredictor.forward"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.predictor.JEPAPredictor.forward" title="Permalink to this definition">#</a></dt>
 <dd></dd></dl>
 
 <dl class="py method">
-<dt class="sig sig-object py" id="cyto_dl.nn.vits.jepa.JEPAPredictor.init_weight">
-<span class="sig-name descname"><span class="pre">init_weight</span></span><span class="sig-paren">(</span><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/jepa.html#JEPAPredictor.init_weight"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.jepa.JEPAPredictor.init_weight" title="Permalink to this definition">#</a></dt>
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.predictor.JEPAPredictor.init_weight">
+<span class="sig-name descname"><span class="pre">init_weight</span></span><span class="sig-paren">(</span><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/predictor.html#JEPAPredictor.init_weight"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.predictor.JEPAPredictor.init_weight" title="Permalink to this definition">#</a></dt>
 <dd></dd></dl>
 
 <dl class="py method">
-<dt class="sig sig-object py" id="cyto_dl.nn.vits.jepa.JEPAPredictor.predict_target_features">
-<span class="sig-name descname"><span class="pre">predict_target_features</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">context_emb</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">target_masks</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/jepa.html#JEPAPredictor.predict_target_features"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.jepa.JEPAPredictor.predict_target_features" title="Permalink to this definition">#</a></dt>
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.predictor.JEPAPredictor.predict_target_features">
+<span class="sig-name descname"><span class="pre">predict_target_features</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">context_emb</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">target_masks</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/predictor.html#JEPAPredictor.predict_target_features"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.predictor.JEPAPredictor.predict_target_features" title="Permalink to this definition">#</a></dt>
 <dd></dd></dl>
 
 </dd></dl>
@@ -526,23 +511,23 @@
       <footer>
         
         <div class="related-pages">
-          <a class="next-page" href="cyto_dl.nn.vits.mae.html">
+          <a class="next-page" href="cyto_dl.nn.vits.seg.html">
               <div class="page-info">
                 <div class="context">
                   <span>Next</span>
                 </div>
-                <div class="title">cyto_dl.nn.vits.mae module</div>
+                <div class="title">cyto_dl.nn.vits.seg module</div>
               </div>
               <svg class="furo-related-icon"><use href="#svg-arrow-right"></use></svg>
             </a>
-          <a class="prev-page" href="cyto_dl.nn.vits.cross_mae.html">
+          <a class="prev-page" href="cyto_dl.nn.vits.mae.html">
               <svg class="furo-related-icon"><use href="#svg-arrow-right"></use></svg>
               <div class="page-info">
                 <div class="context">
                   <span>Previous</span>
                 </div>
                 
-                <div class="title">cyto_dl.nn.vits.cross_mae module</div>
+                <div class="title">cyto_dl.nn.vits.mae module</div>
                 
               </div>
             </a>
@@ -576,19 +561,15 @@
         <div class="toc-tree-container">
           <div class="toc-tree">
             <ul>
-<li><a class="reference internal" href="#">cyto_dl.nn.vits.jepa module</a><ul>
-<li><a class="reference internal" href="#cyto_dl.nn.vits.jepa.IWMPredictor"><code class="docutils literal notranslate"><span class="pre">IWMPredictor</span></code></a><ul>
-<li><a class="reference internal" href="#cyto_dl.nn.vits.jepa.IWMPredictor.forward"><code class="docutils literal notranslate"><span class="pre">IWMPredictor.forward()</span></code></a></li>
-</ul>
-</li>
-<li><a class="reference internal" href="#cyto_dl.nn.vits.jepa.JEPAEncoder"><code class="docutils literal notranslate"><span class="pre">JEPAEncoder</span></code></a><ul>
-<li><a class="reference internal" href="#cyto_dl.nn.vits.jepa.JEPAEncoder.forward"><code class="docutils literal notranslate"><span class="pre">JEPAEncoder.forward()</span></code></a></li>
+<li><a class="reference internal" href="#">cyto_dl.nn.vits.predictor module</a><ul>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.predictor.IWMPredictor"><code class="docutils literal notranslate"><span class="pre">IWMPredictor</span></code></a><ul>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.predictor.IWMPredictor.forward"><code class="docutils literal notranslate"><span class="pre">IWMPredictor.forward()</span></code></a></li>
 </ul>
 </li>
-<li><a class="reference internal" href="#cyto_dl.nn.vits.jepa.JEPAPredictor"><code class="docutils literal notranslate"><span class="pre">JEPAPredictor</span></code></a><ul>
-<li><a class="reference internal" href="#cyto_dl.nn.vits.jepa.JEPAPredictor.forward"><code class="docutils literal notranslate"><span class="pre">JEPAPredictor.forward()</span></code></a></li>
-<li><a class="reference internal" href="#cyto_dl.nn.vits.jepa.JEPAPredictor.init_weight"><code class="docutils literal notranslate"><span class="pre">JEPAPredictor.init_weight()</span></code></a></li>
-<li><a class="reference internal" href="#cyto_dl.nn.vits.jepa.JEPAPredictor.predict_target_features"><code class="docutils literal notranslate"><span class="pre">JEPAPredictor.predict_target_features()</span></code></a></li>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.predictor.JEPAPredictor"><code class="docutils literal notranslate"><span class="pre">JEPAPredictor</span></code></a><ul>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.predictor.JEPAPredictor.forward"><code class="docutils literal notranslate"><span class="pre">JEPAPredictor.forward()</span></code></a></li>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.predictor.JEPAPredictor.init_weight"><code class="docutils literal notranslate"><span class="pre">JEPAPredictor.init_weight()</span></code></a></li>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.predictor.JEPAPredictor.predict_target_features"><code class="docutils literal notranslate"><span class="pre">JEPAPredictor.predict_target_features()</span></code></a></li>
 </ul>
 </li>
 </ul>
diff --git a/cyto_dl.nn.vits.seg.html b/cyto_dl.nn.vits.seg.html
index 1ad7a8b5e..6e1010a5c 100644
--- a/cyto_dl.nn.vits.seg.html
+++ b/cyto_dl.nn.vits.seg.html
@@ -3,7 +3,7 @@
   <head><meta charset="utf-8"/>
     <meta name="viewport" content="width=device-width,initial-scale=1"/>
     <meta name="color-scheme" content="light dark"><meta name="generator" content="Docutils 0.19: https://docutils.sourceforge.io/" />
-<link rel="index" title="Index" href="genindex.html" /><link rel="search" title="Search" href="search.html" /><link rel="next" title="cyto_dl.nn.vits.utils module" href="cyto_dl.nn.vits.utils.html" /><link rel="prev" title="cyto_dl.nn.vits.mae module" href="cyto_dl.nn.vits.mae.html" />
+<link rel="index" title="Index" href="genindex.html" /><link rel="search" title="Search" href="search.html" /><link rel="next" title="cyto_dl.nn.vits.utils module" href="cyto_dl.nn.vits.utils.html" /><link rel="prev" title="cyto_dl.nn.vits.predictor module" href="cyto_dl.nn.vits.predictor.html" />
 
     <!-- Generated with Sphinx 6.1.3 and Furo 2023.03.23 -->
         <title>cyto_dl.nn.vits.seg module - cyto-dl 0.4.0 documentation</title>
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 current has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5 current current-page"><a class="current reference internal" href="#">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
@@ -506,14 +513,14 @@
               </div>
               <svg class="furo-related-icon"><use href="#svg-arrow-right"></use></svg>
             </a>
-          <a class="prev-page" href="cyto_dl.nn.vits.mae.html">
+          <a class="prev-page" href="cyto_dl.nn.vits.predictor.html">
               <svg class="furo-related-icon"><use href="#svg-arrow-right"></use></svg>
               <div class="page-info">
                 <div class="context">
                   <span>Previous</span>
                 </div>
                 
-                <div class="title">cyto_dl.nn.vits.mae module</div>
+                <div class="title">cyto_dl.nn.vits.predictor module</div>
                 
               </div>
             </a>
diff --git a/cyto_dl.nn.vits.utils.html b/cyto_dl.nn.vits.utils.html
index b11561866..c27e003c0 100644
--- a/cyto_dl.nn.vits.utils.html
+++ b/cyto_dl.nn.vits.utils.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 current has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5 current current-page"><a class="current reference internal" href="#">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
@@ -444,6 +451,18 @@
 <p>Can be either learnable or fixed.</p>
 </dd></dl>
 
+<dl class="py function">
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.utils.match_tuple_dimensions">
+<span class="sig-prename descclassname"><span class="pre">cyto_dl.nn.vits.utils.</span></span><span class="sig-name descname"><span class="pre">match_tuple_dimensions</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">spatial_dims</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">tuples</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/utils.html#match_tuple_dimensions"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.utils.match_tuple_dimensions" title="Permalink to this definition">#</a></dt>
+<dd><p>Ensure that each element in a list of tuples has the same length as spatial_dims.</p>
+<p>If a single element, the element is repeated to match the spatial_dims.</p>
+</dd></dl>
+
+<dl class="py function">
+<dt class="sig sig-object py" id="cyto_dl.nn.vits.utils.random_indexes">
+<span class="sig-prename descclassname"><span class="pre">cyto_dl.nn.vits.utils.</span></span><span class="sig-name descname"><span class="pre">random_indexes</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">size</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">int</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">device</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/utils.html#random_indexes"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.utils.random_indexes" title="Permalink to this definition">#</a></dt>
+<dd></dd></dl>
+
 <dl class="py function">
 <dt class="sig sig-object py" id="cyto_dl.nn.vits.utils.take_indexes">
 <span class="sig-prename descclassname"><span class="pre">cyto_dl.nn.vits.utils.</span></span><span class="sig-name descname"><span class="pre">take_indexes</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">sequences</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">indexes</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/cyto_dl/nn/vits/utils.html#take_indexes"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#cyto_dl.nn.vits.utils.take_indexes" title="Permalink to this definition">#</a></dt>
@@ -508,6 +527,8 @@
             <ul>
 <li><a class="reference internal" href="#">cyto_dl.nn.vits.utils module</a><ul>
 <li><a class="reference internal" href="#cyto_dl.nn.vits.utils.get_positional_embedding"><code class="docutils literal notranslate"><span class="pre">get_positional_embedding()</span></code></a></li>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.utils.match_tuple_dimensions"><code class="docutils literal notranslate"><span class="pre">match_tuple_dimensions()</span></code></a></li>
+<li><a class="reference internal" href="#cyto_dl.nn.vits.utils.random_indexes"><code class="docutils literal notranslate"><span class="pre">random_indexes()</span></code></a></li>
 <li><a class="reference internal" href="#cyto_dl.nn.vits.utils.take_indexes"><code class="docutils literal notranslate"><span class="pre">take_indexes()</span></code></a></li>
 </ul>
 </li>
diff --git a/cyto_dl.point_cloud.html b/cyto_dl.point_cloud.html
index 6a2a69924..ea7f1e3d5 100644
--- a/cyto_dl.point_cloud.html
+++ b/cyto_dl.point_cloud.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 current has-children current-page"><a class="current reference internal" href="#">cyto_dl.point_cloud package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 current has-children current-page"><a class="current reference internal" href="#">cyto_dl.point_cloud package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.point_cloud.io.html b/cyto_dl.point_cloud.io.html
index 5e02b2bae..eb2285634 100644
--- a/cyto_dl.point_cloud.io.html
+++ b/cyto_dl.point_cloud.io.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 current has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
-<li class="toctree-l4 current has-children current-page"><a class="current reference internal" href="#">cyto_dl.point_cloud.io package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 current has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
+<li class="toctree-l4 current has-children current-page"><a class="current reference internal" href="#">cyto_dl.point_cloud.io package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.point_cloud.io.read_pcloud.html b/cyto_dl.point_cloud.io.read_pcloud.html
index 6033dd0f6..e3b6f7bbe 100644
--- a/cyto_dl.point_cloud.io.read_pcloud.html
+++ b/cyto_dl.point_cloud.io.read_pcloud.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 current has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
-<li class="toctree-l4 current has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
+<li class="toctree-l3 current has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
+<li class="toctree-l4 current has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
 <li class="toctree-l5 current current-page"><a class="current reference internal" href="#">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.train.html b/cyto_dl.train.html
index 180233800..aee5e6b0a 100644
--- a/cyto_dl.train.html
+++ b/cyto_dl.train.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.utils.array.html b/cyto_dl.utils.array.html
index 248d8e1ae..94f61fb80 100644
--- a/cyto_dl.utils.array.html
+++ b/cyto_dl.utils.array.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 current has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 current has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.utils.checkpoint.html b/cyto_dl.utils.checkpoint.html
index 50a53139d..f6ff17a4b 100644
--- a/cyto_dl.utils.checkpoint.html
+++ b/cyto_dl.utils.checkpoint.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 current has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 current has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.utils.config.html b/cyto_dl.utils.config.html
index c92902f68..22e791875 100644
--- a/cyto_dl.utils.config.html
+++ b/cyto_dl.utils.config.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 current has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 current has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.utils.download_test_data.html b/cyto_dl.utils.download_test_data.html
index d2dc78ceb..a8f9fb8a0 100644
--- a/cyto_dl.utils.download_test_data.html
+++ b/cyto_dl.utils.download_test_data.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 current has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 current has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.utils.dummy_dataset.html b/cyto_dl.utils.dummy_dataset.html
index 35fd14f2c..71319dd90 100644
--- a/cyto_dl.utils.dummy_dataset.html
+++ b/cyto_dl.utils.dummy_dataset.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 current has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 current has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.utils.html b/cyto_dl.utils.html
index 11e9916db..8c12ad01f 100644
--- a/cyto_dl.utils.html
+++ b/cyto_dl.utils.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 current has-children current-page"><a class="current reference internal" href="#">cyto_dl.utils package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 current has-children current-page"><a class="current reference internal" href="#">cyto_dl.utils package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.utils.pylogger.html b/cyto_dl.utils.pylogger.html
index 31d4ccef3..c19305c9d 100644
--- a/cyto_dl.utils.pylogger.html
+++ b/cyto_dl.utils.pylogger.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 current has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 current has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.utils.rich_utils.html b/cyto_dl.utils.rich_utils.html
index de54df066..452477630 100644
--- a/cyto_dl.utils.rich_utils.html
+++ b/cyto_dl.utils.rich_utils.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 current has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 current has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.utils.rotation.html b/cyto_dl.utils.rotation.html
index 5c05d54ea..4e2106e6e 100644
--- a/cyto_dl.utils.rotation.html
+++ b/cyto_dl.utils.rotation.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 current has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 current has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.utils.spharm.html b/cyto_dl.utils.spharm.html
index 283446895..2339bffef 100644
--- a/cyto_dl.utils.spharm.html
+++ b/cyto_dl.utils.spharm.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 current has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
-<li class="toctree-l4 current has-children current-page"><a class="current reference internal" href="#">cyto_dl.utils.spharm package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 current has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
+<li class="toctree-l4 current has-children current-page"><a class="current reference internal" href="#">cyto_dl.utils.spharm package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.utils.spharm.reconstruction.html b/cyto_dl.utils.spharm.reconstruction.html
index 5540dd7fc..2c263494b 100644
--- a/cyto_dl.utils.spharm.reconstruction.html
+++ b/cyto_dl.utils.spharm.reconstruction.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 current has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
-<li class="toctree-l4 current has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
+<li class="toctree-l3 current has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
+<li class="toctree-l4 current has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
 <li class="toctree-l5 current current-page"><a class="current reference internal" href="#">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.utils.spharm.rotation.html b/cyto_dl.utils.spharm.rotation.html
index 16e045da6..bbf8070aa 100644
--- a/cyto_dl.utils.spharm.rotation.html
+++ b/cyto_dl.utils.spharm.rotation.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 current has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
-<li class="toctree-l4 current has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
+<li class="toctree-l3 current has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
+<li class="toctree-l4 current has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5 current current-page"><a class="current reference internal" href="#">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/cyto_dl.utils.template_utils.html b/cyto_dl.utils.template_utils.html
index ff5c69043..baf7f6478 100644
--- a/cyto_dl.utils.template_utils.html
+++ b/cyto_dl.utils.template_utils.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 current has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 current has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input checked="" class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul class="current">
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/genindex.html b/genindex.html
index dbd667844..4db741f9d 100644
--- a/genindex.html
+++ b/genindex.html
@@ -348,14 +348,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -367,15 +374,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
@@ -441,17 +448,21 @@ <h2>A</h2>
   <table style="width: 100%" class="indextable genindextable"><tr>
     <td style="width: 33%; vertical-align: top;"><ul>
         <li><a href="cyto_dl.models.im2im.utils.postprocessing.act_thresh_label.html#cyto_dl.models.im2im.utils.postprocessing.act_thresh_label.ActThreshLabel">ActThreshLabel (class in cyto_dl.models.im2im.utils.postprocessing.act_thresh_label)</a>
+</li>
+        <li><a href="cyto_dl.nn.vits.decoder.html#cyto_dl.nn.vits.decoder.MAE_Decoder.add_mask_tokens">add_mask_tokens() (cyto_dl.nn.vits.decoder.MAE_Decoder method)</a>
 </li>
         <li><a href="cyto_dl.image.transforms.add_meta.html#cyto_dl.image.transforms.add_meta.AddMeta">AddMeta (class in cyto_dl.image.transforms.add_meta)</a>
+</li>
+        <li><a href="cyto_dl.nn.vits.decoder.html#cyto_dl.nn.vits.decoder.MAE_Decoder.adjust_indices_for_cls">adjust_indices_for_cls() (cyto_dl.nn.vits.decoder.MAE_Decoder method)</a>
 </li>
         <li><a href="cyto_dl.nn.losses.adversarial_loss.html#cyto_dl.nn.losses.adversarial_loss.AdversarialLoss">AdversarialLoss (class in cyto_dl.nn.losses.adversarial_loss)</a>
 </li>
         <li><a href="cyto_dl.loggers.mlflow.html#cyto_dl.loggers.mlflow.MLFlowLogger.after_save_checkpoint">after_save_checkpoint() (cyto_dl.loggers.mlflow.MLFlowLogger method)</a>
-</li>
-        <li><a href="cyto_dl.image.io.aicsimage_loader.html#cyto_dl.image.io.aicsimage_loader.AICSImageLoaderd">AICSImageLoaderd (class in cyto_dl.image.io.aicsimage_loader)</a>
 </li>
     </ul></td>
     <td style="width: 33%; vertical-align: top;"><ul>
+        <li><a href="cyto_dl.image.io.aicsimage_loader.html#cyto_dl.image.io.aicsimage_loader.AICSImageLoaderd">AICSImageLoaderd (class in cyto_dl.image.io.aicsimage_loader)</a>
+</li>
         <li><a href="cyto_dl.datamodules.dataframe.utils.html#cyto_dl.datamodules.dataframe.utils.AlternatingBatchSampler">AlternatingBatchSampler (class in cyto_dl.datamodules.dataframe.utils)</a>
 </li>
         <li><a href="cyto_dl.dataframe.transforms.misc.html#cyto_dl.dataframe.transforms.misc.append_class_weights">append_class_weights() (in module cyto_dl.dataframe.transforms.misc)</a>
@@ -569,10 +580,20 @@ <h2>C</h2>
         <li><a href="cyto_dl.datamodules.torchvision.html#cyto_dl.datamodules.torchvision.Country211">Country211 (class in cyto_dl.datamodules.torchvision)</a>
 </li>
         <li><a href="cyto_dl.api.data.html#cyto_dl.api.data.HardwareType.CPU">CPU (cyto_dl.api.data.HardwareType attribute)</a>
+</li>
+        <li><a href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html#cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase.create_conv">create_conv() (cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase method)</a>
 </li>
         <li><a href="cyto_dl.utils.array.html#cyto_dl.utils.array.create_dataloader">create_dataloader() (in module cyto_dl.utils.array)</a>
 </li>
+        <li><a href="cyto_dl.nn.vits.blocks.patchify.patchify.html#cyto_dl.nn.vits.blocks.patchify.patchify.Patchify.create_img2token">create_img2token() (cyto_dl.nn.vits.blocks.patchify.patchify.Patchify method)</a>
+
+        <ul>
+          <li><a href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html#cyto_dl.nn.vits.blocks.patchify.patchify_hiera.PatchifyHiera.create_img2token">(cyto_dl.nn.vits.blocks.patchify.patchify_hiera.PatchifyHiera method)</a>
+</li>
+        </ul></li>
         <li><a href="cyto_dl.datamodules.multidim_image.html#cyto_dl.datamodules.multidim_image.MultiDimImageDataset.create_metatensor">create_metatensor() (cyto_dl.datamodules.multidim_image.MultiDimImageDataset method)</a>
+</li>
+        <li><a href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html#cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase.create_patch2img">create_patch2img() (cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase method)</a>
 </li>
         <li><a href="cyto_dl.image.transforms.track_transforms.html#cyto_dl.image.transforms.track_transforms.CropResize">CropResize (class in cyto_dl.image.transforms.track_transforms)</a>
 </li>
@@ -580,7 +601,7 @@ <h2>C</h2>
 </li>
         <li><a href="cyto_dl.nn.vits.blocks.cross_attention.html#cyto_dl.nn.vits.blocks.cross_attention.CrossAttentionBlock">CrossAttentionBlock (class in cyto_dl.nn.vits.blocks.cross_attention)</a>
 </li>
-        <li><a href="cyto_dl.nn.vits.cross_mae.html#cyto_dl.nn.vits.cross_mae.CrossMAE_Decoder">CrossMAE_Decoder (class in cyto_dl.nn.vits.cross_mae)</a>
+        <li><a href="cyto_dl.nn.vits.decoder.html#cyto_dl.nn.vits.decoder.CrossMAE_Decoder">CrossMAE_Decoder (class in cyto_dl.nn.vits.decoder)</a>
 </li>
         <li><a href="cyto_dl.nn.vits.blocks.cross_attention.html#cyto_dl.nn.vits.blocks.cross_attention.CrossSelfBlock">CrossSelfBlock (class in cyto_dl.nn.vits.blocks.cross_attention)</a>
 </li>
@@ -992,8 +1013,6 @@ <h2>C</h2>
           <li><a href="cyto_dl.models.classification.html#module-cyto_dl.models.classification">module</a>
 </li>
         </ul></li>
-    </ul></td>
-    <td style="width: 33%; vertical-align: top;"><ul>
         <li>
     cyto_dl.models.classification.classification
 
@@ -1001,6 +1020,8 @@ <h2>C</h2>
           <li><a href="cyto_dl.models.classification.classification.html#module-cyto_dl.models.classification.classification">module</a>
 </li>
         </ul></li>
+    </ul></td>
+    <td style="width: 33%; vertical-align: top;"><ul>
         <li>
     cyto_dl.models.classification.timepoint_classification
 
@@ -1412,6 +1433,13 @@ <h2>C</h2>
 
         <ul>
           <li><a href="cyto_dl.nn.vits.blocks.intermediate_weigher.html#module-cyto_dl.nn.vits.blocks.intermediate_weigher">module</a>
+</li>
+        </ul></li>
+        <li>
+    cyto_dl.nn.vits.blocks.masked_unit_attention
+
+        <ul>
+          <li><a href="cyto_dl.nn.vits.blocks.masked_unit_attention.html#module-cyto_dl.nn.vits.blocks.masked_unit_attention">module</a>
 </li>
         </ul></li>
         <li>
@@ -1422,17 +1450,38 @@ <h2>C</h2>
 </li>
         </ul></li>
         <li>
-    cyto_dl.nn.vits.cross_mae
+    cyto_dl.nn.vits.blocks.patchify.patchify
+
+        <ul>
+          <li><a href="cyto_dl.nn.vits.blocks.patchify.patchify.html#module-cyto_dl.nn.vits.blocks.patchify.patchify">module</a>
+</li>
+        </ul></li>
+        <li>
+    cyto_dl.nn.vits.blocks.patchify.patchify_base
 
         <ul>
-          <li><a href="cyto_dl.nn.vits.cross_mae.html#module-cyto_dl.nn.vits.cross_mae">module</a>
+          <li><a href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html#module-cyto_dl.nn.vits.blocks.patchify.patchify_base">module</a>
 </li>
         </ul></li>
         <li>
-    cyto_dl.nn.vits.jepa
+    cyto_dl.nn.vits.blocks.patchify.patchify_hiera
 
         <ul>
-          <li><a href="cyto_dl.nn.vits.jepa.html#module-cyto_dl.nn.vits.jepa">module</a>
+          <li><a href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html#module-cyto_dl.nn.vits.blocks.patchify.patchify_hiera">module</a>
+</li>
+        </ul></li>
+        <li>
+    cyto_dl.nn.vits.decoder
+
+        <ul>
+          <li><a href="cyto_dl.nn.vits.decoder.html#module-cyto_dl.nn.vits.decoder">module</a>
+</li>
+        </ul></li>
+        <li>
+    cyto_dl.nn.vits.encoder
+
+        <ul>
+          <li><a href="cyto_dl.nn.vits.encoder.html#module-cyto_dl.nn.vits.encoder">module</a>
 </li>
         </ul></li>
         <li>
@@ -1440,6 +1489,13 @@ <h2>C</h2>
 
         <ul>
           <li><a href="cyto_dl.nn.vits.mae.html#module-cyto_dl.nn.vits.mae">module</a>
+</li>
+        </ul></li>
+        <li>
+    cyto_dl.nn.vits.predictor
+
+        <ul>
+          <li><a href="cyto_dl.nn.vits.predictor.html#module-cyto_dl.nn.vits.predictor">module</a>
 </li>
         </ul></li>
         <li>
@@ -1558,14 +1614,22 @@ <h2>D</h2>
 </li>
         <li><a href="cyto_dl.models.vae.base_vae.html#cyto_dl.models.vae.base_vae.BaseVAE.decode">decode() (cyto_dl.models.vae.base_vae.BaseVAE method)</a>
 </li>
+        <li><a href="cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.HieraMAE.decoder">decoder (cyto_dl.nn.vits.mae.HieraMAE property)</a>
+
+        <ul>
+          <li><a href="cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.MAE.decoder">(cyto_dl.nn.vits.mae.MAE property)</a>
+</li>
+          <li><a href="cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.MAE_Base.decoder">(cyto_dl.nn.vits.mae.MAE_Base property)</a>
+</li>
+        </ul></li>
         <li><a href="cyto_dl.utils.download_test_data.html#cyto_dl.utils.download_test_data.delete_test_data">delete_test_data() (in module cyto_dl.utils.download_test_data)</a>
 </li>
+    </ul></td>
+    <td style="width: 33%; vertical-align: top;"><ul>
         <li><a href="cyto_dl.models.im2im.utils.postprocessing.dict_to_im.html#cyto_dl.models.im2im.utils.postprocessing.dict_to_im.detach">detach() (in module cyto_dl.models.im2im.utils.postprocessing.dict_to_im)</a>
 </li>
         <li><a href="cyto_dl.models.vae.priors.gaussian.html#cyto_dl.models.vae.priors.gaussian.DiagonalGaussianPrior">DiagonalGaussianPrior (class in cyto_dl.models.vae.priors.gaussian)</a>
 </li>
-    </ul></td>
-    <td style="width: 33%; vertical-align: top;"><ul>
         <li><a href="cyto_dl.models.im2im.utils.postprocessing.dict_to_im.html#cyto_dl.models.im2im.utils.postprocessing.dict_to_im.DictToIm">DictToIm (class in cyto_dl.models.im2im.utils.postprocessing.dict_to_im)</a>
 </li>
         <li><a href="cyto_dl.api.model.html#cyto_dl.api.model.CytoDLModel.download_example_data">download_example_data() (cyto_dl.api.model.CytoDLModel method)</a>
@@ -1594,6 +1658,14 @@ <h2>E</h2>
 </li>
         <li><a href="cyto_dl.nn.vits.seg.html#cyto_dl.nn.vits.seg.EncodedSkip">EncodedSkip (class in cyto_dl.nn.vits.seg)</a>
 </li>
+        <li><a href="cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.HieraMAE.encoder">encoder (cyto_dl.nn.vits.mae.HieraMAE property)</a>
+
+        <ul>
+          <li><a href="cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.MAE.encoder">(cyto_dl.nn.vits.mae.MAE property)</a>
+</li>
+          <li><a href="cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.MAE_Base.encoder">(cyto_dl.nn.vits.mae.MAE_Base property)</a>
+</li>
+        </ul></li>
         <li><a href="cyto_dl.utils.rich_utils.html#cyto_dl.utils.rich_utils.enforce_tags">enforce_tags() (in module cyto_dl.utils.rich_utils)</a>
 </li>
     </ul></td>
@@ -1606,6 +1678,14 @@ <h2>E</h2>
 </li>
         <li><a href="cyto_dl.models.jepa.iwm.html#cyto_dl.models.jepa.iwm.IWM.extract_embeddings">extract_embeddings() (cyto_dl.models.jepa.iwm.IWM method)</a>
 </li>
+        <li><a href="cyto_dl.nn.vits.blocks.patchify.patchify.html#cyto_dl.nn.vits.blocks.patchify.patchify.Patchify.extract_visible_tokens">extract_visible_tokens() (cyto_dl.nn.vits.blocks.patchify.patchify.Patchify method)</a>
+
+        <ul>
+          <li><a href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html#cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase.extract_visible_tokens">(cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase method)</a>
+</li>
+          <li><a href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html#cyto_dl.nn.vits.blocks.patchify.patchify_hiera.PatchifyHiera.extract_visible_tokens">(cyto_dl.nn.vits.blocks.patchify.patchify_hiera.PatchifyHiera method)</a>
+</li>
+        </ul></li>
         <li><a href="cyto_dl.utils.template_utils.html#cyto_dl.utils.template_utils.extras">extras() (in module cyto_dl.utils.template_utils)</a>
 </li>
     </ul></td>
@@ -1719,21 +1799,29 @@ <h2>F</h2>
 </li>
           <li><a href="cyto_dl.nn.vits.blocks.intermediate_weigher.html#cyto_dl.nn.vits.blocks.intermediate_weigher.IntermediateWeigher.forward">(cyto_dl.nn.vits.blocks.intermediate_weigher.IntermediateWeigher method)</a>
 </li>
-          <li><a href="cyto_dl.nn.vits.blocks.patchify.html#cyto_dl.nn.vits.blocks.patchify.Patchify.forward">(cyto_dl.nn.vits.blocks.patchify.Patchify method)</a>
+          <li><a href="cyto_dl.nn.vits.blocks.masked_unit_attention.html#cyto_dl.nn.vits.blocks.masked_unit_attention.HieraBlock.forward">(cyto_dl.nn.vits.blocks.masked_unit_attention.HieraBlock method)</a>
 </li>
-          <li><a href="cyto_dl.nn.vits.cross_mae.html#cyto_dl.nn.vits.cross_mae.CrossMAE_Decoder.forward">(cyto_dl.nn.vits.cross_mae.CrossMAE_Decoder method)</a>
+          <li><a href="cyto_dl.nn.vits.blocks.masked_unit_attention.html#cyto_dl.nn.vits.blocks.masked_unit_attention.MaskUnitAttention.forward">(cyto_dl.nn.vits.blocks.masked_unit_attention.MaskUnitAttention method)</a>
 </li>
-          <li><a href="cyto_dl.nn.vits.jepa.html#cyto_dl.nn.vits.jepa.IWMPredictor.forward">(cyto_dl.nn.vits.jepa.IWMPredictor method)</a>
+          <li><a href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html#cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase.forward">(cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase method)</a>
 </li>
-          <li><a href="cyto_dl.nn.vits.jepa.html#cyto_dl.nn.vits.jepa.JEPAEncoder.forward">(cyto_dl.nn.vits.jepa.JEPAEncoder method)</a>
+          <li><a href="cyto_dl.nn.vits.decoder.html#cyto_dl.nn.vits.decoder.CrossMAE_Decoder.forward">(cyto_dl.nn.vits.decoder.CrossMAE_Decoder method)</a>
 </li>
-          <li><a href="cyto_dl.nn.vits.jepa.html#cyto_dl.nn.vits.jepa.JEPAPredictor.forward">(cyto_dl.nn.vits.jepa.JEPAPredictor method)</a>
+          <li><a href="cyto_dl.nn.vits.decoder.html#cyto_dl.nn.vits.decoder.MAE_Decoder.forward">(cyto_dl.nn.vits.decoder.MAE_Decoder method)</a>
 </li>
-          <li><a href="cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.MAE_Decoder.forward">(cyto_dl.nn.vits.mae.MAE_Decoder method)</a>
+          <li><a href="cyto_dl.nn.vits.encoder.html#cyto_dl.nn.vits.encoder.HieraEncoder.forward">(cyto_dl.nn.vits.encoder.HieraEncoder method)</a>
 </li>
-          <li><a href="cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.MAE_Encoder.forward">(cyto_dl.nn.vits.mae.MAE_Encoder method)</a>
+          <li><a href="cyto_dl.nn.vits.encoder.html#cyto_dl.nn.vits.encoder.JEPAEncoder.forward">(cyto_dl.nn.vits.encoder.JEPAEncoder method)</a>
 </li>
-          <li><a href="cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.MAE_ViT.forward">(cyto_dl.nn.vits.mae.MAE_ViT method)</a>
+          <li><a href="cyto_dl.nn.vits.encoder.html#cyto_dl.nn.vits.encoder.MAE_Encoder.forward">(cyto_dl.nn.vits.encoder.MAE_Encoder method)</a>
+</li>
+          <li><a href="cyto_dl.nn.vits.encoder.html#cyto_dl.nn.vits.encoder.SpatialMerger.forward">(cyto_dl.nn.vits.encoder.SpatialMerger method)</a>
+</li>
+          <li><a href="cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.MAE_Base.forward">(cyto_dl.nn.vits.mae.MAE_Base method)</a>
+</li>
+          <li><a href="cyto_dl.nn.vits.predictor.html#cyto_dl.nn.vits.predictor.IWMPredictor.forward">(cyto_dl.nn.vits.predictor.IWMPredictor method)</a>
+</li>
+          <li><a href="cyto_dl.nn.vits.predictor.html#cyto_dl.nn.vits.predictor.JEPAPredictor.forward">(cyto_dl.nn.vits.predictor.JEPAPredictor method)</a>
 </li>
           <li><a href="cyto_dl.nn.vits.seg.html#cyto_dl.nn.vits.seg.EncodedSkip.forward">(cyto_dl.nn.vits.seg.EncodedSkip method)</a>
 </li>
@@ -1816,10 +1904,10 @@ <h2>G</h2>
 </li>
         <li><a href="cyto_dl.nn.losses.gan_loss.html#cyto_dl.nn.losses.gan_loss.Pix2PixHD.get_gan_loss">get_gan_loss() (cyto_dl.nn.losses.gan_loss.Pix2PixHD method)</a>
 </li>
-    </ul></td>
-    <td style="width: 33%; vertical-align: top;"><ul>
         <li><a href="cyto_dl.api.cyto_dl_model.segmentation_plugin_model.html#cyto_dl.api.cyto_dl_model.segmentation_plugin_model.SegmentationPluginModel.get_hardware_type">get_hardware_type() (cyto_dl.api.cyto_dl_model.segmentation_plugin_model.SegmentationPluginModel method)</a>
 </li>
+    </ul></td>
+    <td style="width: 33%; vertical-align: top;"><ul>
         <li><a href="cyto_dl.api.cyto_dl_model.segmentation_plugin_model.html#cyto_dl.api.cyto_dl_model.segmentation_plugin_model.SegmentationPluginModel.get_input_channel">get_input_channel() (cyto_dl.api.cyto_dl_model.segmentation_plugin_model.SegmentationPluginModel method)</a>
 </li>
         <li><a href="cyto_dl.api.cyto_dl_model.segmentation_plugin_model.html#cyto_dl.api.cyto_dl_model.segmentation_plugin_model.SegmentationPluginModel.get_manifest_column_names">get_manifest_column_names() (cyto_dl.api.cyto_dl_model.segmentation_plugin_model.SegmentationPluginModel method)</a>
@@ -1827,7 +1915,15 @@ <h2>G</h2>
         <li><a href="cyto_dl.models.jepa.jepa_base.html#cyto_dl.models.jepa.jepa_base.JEPABase.get_mask">get_mask() (cyto_dl.models.jepa.jepa_base.JEPABase method)</a>
 
         <ul>
-          <li><a href="cyto_dl.nn.vits.blocks.patchify.html#cyto_dl.nn.vits.blocks.patchify.Patchify.get_mask">(cyto_dl.nn.vits.blocks.patchify.Patchify method)</a>
+          <li><a href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html#cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase.get_mask">(cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase method)</a>
+</li>
+        </ul></li>
+        <li><a href="cyto_dl.nn.vits.blocks.patchify.patchify.html#cyto_dl.nn.vits.blocks.patchify.patchify.Patchify.get_mask_args">get_mask_args() (cyto_dl.nn.vits.blocks.patchify.patchify.Patchify method)</a>
+
+        <ul>
+          <li><a href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html#cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase.get_mask_args">(cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase method)</a>
+</li>
+          <li><a href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html#cyto_dl.nn.vits.blocks.patchify.patchify_hiera.PatchifyHiera.get_mask_args">(cyto_dl.nn.vits.blocks.patchify.patchify_hiera.PatchifyHiera method)</a>
 </li>
         </ul></li>
         <li><a href="cyto_dl.utils.template_utils.html#cyto_dl.utils.template_utils.get_metric_value">get_metric_value() (in module cyto_dl.utils.template_utils)</a>
@@ -1881,9 +1977,15 @@ <h2>H</h2>
   <table style="width: 100%" class="indextable genindextable"><tr>
     <td style="width: 33%; vertical-align: top;"><ul>
         <li><a href="cyto_dl.api.data.html#cyto_dl.api.data.HardwareType">HardwareType (class in cyto_dl.api.data)</a>
+</li>
+        <li><a href="cyto_dl.nn.vits.blocks.masked_unit_attention.html#cyto_dl.nn.vits.blocks.masked_unit_attention.HieraBlock">HieraBlock (class in cyto_dl.nn.vits.blocks.masked_unit_attention)</a>
+</li>
+        <li><a href="cyto_dl.nn.vits.encoder.html#cyto_dl.nn.vits.encoder.HieraEncoder">HieraEncoder (class in cyto_dl.nn.vits.encoder)</a>
 </li>
     </ul></td>
     <td style="width: 33%; vertical-align: top;"><ul>
+        <li><a href="cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.HieraMAE">HieraMAE (class in cyto_dl.nn.vits.mae)</a>
+</li>
         <li><a href="cyto_dl.datamodules.torchvision.html#cyto_dl.datamodules.torchvision.HMDB51">HMDB51 (class in cyto_dl.datamodules.torchvision)</a>
 </li>
         <li><a href="cyto_dl.nn.hr_skip.html#cyto_dl.nn.hr_skip.HRSkip">HRSkip (class in cyto_dl.nn.hr_skip)</a>
@@ -1906,24 +2008,34 @@ <h2>I</h2>
 </li>
         <li><a href="cyto_dl.callbacks.image_saver.html#cyto_dl.callbacks.image_saver.ImageSaver">ImageSaver (class in cyto_dl.callbacks.image_saver)</a>
 </li>
+        <li><a href="cyto_dl.nn.vits.blocks.patchify.patchify.html#cyto_dl.nn.vits.blocks.patchify.patchify.Patchify.img2token">img2token (cyto_dl.nn.vits.blocks.patchify.patchify.Patchify property)</a>
+
+        <ul>
+          <li><a href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html#cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase.img2token">(cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase property)</a>
+</li>
+          <li><a href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html#cyto_dl.nn.vits.blocks.patchify.patchify_hiera.PatchifyHiera.img2token">(cyto_dl.nn.vits.blocks.patchify.patchify_hiera.PatchifyHiera property)</a>
+</li>
+        </ul></li>
         <li><a href="cyto_dl.datamodules.torchvision.html#cyto_dl.datamodules.torchvision.INaturalist">INaturalist (class in cyto_dl.datamodules.torchvision)</a>
 </li>
-        <li><a href="cyto_dl.nn.vits.cross_mae.html#cyto_dl.nn.vits.cross_mae.CrossMAE_Decoder.init_weight">init_weight() (cyto_dl.nn.vits.cross_mae.CrossMAE_Decoder method)</a>
+        <li><a href="cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.MAE_Base.init_decoder">init_decoder() (cyto_dl.nn.vits.mae.MAE_Base method)</a>
+</li>
+        <li><a href="cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.MAE_Base.init_encoder">init_encoder() (cyto_dl.nn.vits.mae.MAE_Base method)</a>
+</li>
+        <li><a href="cyto_dl.nn.vits.decoder.html#cyto_dl.nn.vits.decoder.MAE_Decoder.init_weight">init_weight() (cyto_dl.nn.vits.decoder.MAE_Decoder method)</a>
 
         <ul>
-          <li><a href="cyto_dl.nn.vits.jepa.html#cyto_dl.nn.vits.jepa.JEPAPredictor.init_weight">(cyto_dl.nn.vits.jepa.JEPAPredictor method)</a>
+          <li><a href="cyto_dl.nn.vits.encoder.html#cyto_dl.nn.vits.encoder.MAE_Encoder.init_weight">(cyto_dl.nn.vits.encoder.MAE_Encoder method)</a>
 </li>
-          <li><a href="cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.MAE_Decoder.init_weight">(cyto_dl.nn.vits.mae.MAE_Decoder method)</a>
-</li>
-          <li><a href="cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.MAE_Encoder.init_weight">(cyto_dl.nn.vits.mae.MAE_Encoder method)</a>
+          <li><a href="cyto_dl.nn.vits.predictor.html#cyto_dl.nn.vits.predictor.JEPAPredictor.init_weight">(cyto_dl.nn.vits.predictor.JEPAPredictor method)</a>
 </li>
         </ul></li>
+    </ul></td>
+    <td style="width: 33%; vertical-align: top;"><ul>
         <li><a href="cyto_dl.api.data.html#cyto_dl.api.data.ExperimentType.INSTANCE_SEG">INSTANCE_SEG (cyto_dl.api.data.ExperimentType attribute)</a>
 </li>
         <li><a href="cyto_dl.models.im2im.utils.instance_seg.html#cyto_dl.models.im2im.utils.instance_seg.InstanceSegCluster">InstanceSegCluster (class in cyto_dl.models.im2im.utils.instance_seg)</a>
 </li>
-    </ul></td>
-    <td style="width: 33%; vertical-align: top;"><ul>
         <li><a href="cyto_dl.models.im2im.utils.instance_seg.html#cyto_dl.models.im2im.utils.instance_seg.InstanceSegLoss">InstanceSegLoss (class in cyto_dl.models.im2im.utils.instance_seg)</a>
 </li>
         <li><a href="cyto_dl.models.im2im.utils.instance_seg.html#cyto_dl.models.im2im.utils.instance_seg.InstanceSegPreprocessd">InstanceSegPreprocessd (class in cyto_dl.models.im2im.utils.instance_seg)</a>
@@ -1946,7 +2058,7 @@ <h2>I</h2>
 </li>
         <li><a href="cyto_dl.models.jepa.iwm.html#cyto_dl.models.jepa.iwm.IWM">IWM (class in cyto_dl.models.jepa.iwm)</a>
 </li>
-        <li><a href="cyto_dl.nn.vits.jepa.html#cyto_dl.nn.vits.jepa.IWMPredictor">IWMPredictor (class in cyto_dl.nn.vits.jepa)</a>
+        <li><a href="cyto_dl.nn.vits.predictor.html#cyto_dl.nn.vits.predictor.IWMPredictor">IWMPredictor (class in cyto_dl.nn.vits.predictor)</a>
 </li>
     </ul></td>
   </tr></table>
@@ -1958,13 +2070,13 @@ <h2>J</h2>
     <td style="width: 33%; vertical-align: top;"><ul>
         <li><a href="cyto_dl.models.jepa.jepa_base.html#cyto_dl.models.jepa.jepa_base.JEPABase">JEPABase (class in cyto_dl.models.jepa.jepa_base)</a>
 </li>
-        <li><a href="cyto_dl.nn.vits.jepa.html#cyto_dl.nn.vits.jepa.JEPAEncoder">JEPAEncoder (class in cyto_dl.nn.vits.jepa)</a>
+        <li><a href="cyto_dl.nn.vits.encoder.html#cyto_dl.nn.vits.encoder.JEPAEncoder">JEPAEncoder (class in cyto_dl.nn.vits.encoder)</a>
 </li>
         <li><a href="cyto_dl.image.transforms.generate_jepa_masks.html#cyto_dl.image.transforms.generate_jepa_masks.JEPAMaskGenerator">JEPAMaskGenerator (class in cyto_dl.image.transforms.generate_jepa_masks)</a>
 </li>
     </ul></td>
     <td style="width: 33%; vertical-align: top;"><ul>
-        <li><a href="cyto_dl.nn.vits.jepa.html#cyto_dl.nn.vits.jepa.JEPAPredictor">JEPAPredictor (class in cyto_dl.nn.vits.jepa)</a>
+        <li><a href="cyto_dl.nn.vits.predictor.html#cyto_dl.nn.vits.predictor.JEPAPredictor">JEPAPredictor (class in cyto_dl.nn.vits.predictor)</a>
 </li>
         <li><a href="cyto_dl.callbacks.csv_saver.html#cyto_dl.callbacks.csv_saver.JEPASaver">JEPASaver (class in cyto_dl.callbacks.csv_saver)</a>
 </li>
@@ -2058,11 +2170,13 @@ <h2>L</h2>
   <h2>M</h2>
   <table style="width: 100%" class="indextable genindextable"><tr>
     <td style="width: 33%; vertical-align: top;"><ul>
-        <li><a href="cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.MAE_Decoder">MAE_Decoder (class in cyto_dl.nn.vits.mae)</a>
+        <li><a href="cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.MAE">MAE (class in cyto_dl.nn.vits.mae)</a>
 </li>
-        <li><a href="cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.MAE_Encoder">MAE_Encoder (class in cyto_dl.nn.vits.mae)</a>
+        <li><a href="cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.MAE_Base">MAE_Base (class in cyto_dl.nn.vits.mae)</a>
 </li>
-        <li><a href="cyto_dl.nn.vits.mae.html#cyto_dl.nn.vits.mae.MAE_ViT">MAE_ViT (class in cyto_dl.nn.vits.mae)</a>
+        <li><a href="cyto_dl.nn.vits.decoder.html#cyto_dl.nn.vits.decoder.MAE_Decoder">MAE_Decoder (class in cyto_dl.nn.vits.decoder)</a>
+</li>
+        <li><a href="cyto_dl.nn.vits.encoder.html#cyto_dl.nn.vits.encoder.MAE_Encoder">MAE_Encoder (class in cyto_dl.nn.vits.encoder)</a>
 </li>
         <li><a href="cyto_dl.nn.head.mae_head.html#cyto_dl.nn.head.mae_head.MAEHead">MAEHead (class in cyto_dl.nn.head.mae_head)</a>
 </li>
@@ -2095,6 +2209,10 @@ <h2>M</h2>
         <li><a href="cyto_dl.datamodules.dataframe.utils.html#cyto_dl.datamodules.dataframe.utils.make_single_dataframe_splits">make_single_dataframe_splits() (in module cyto_dl.datamodules.dataframe.utils)</a>
 </li>
         <li><a href="cyto_dl.nn.head.mask_head.html#cyto_dl.nn.head.mask_head.MaskHead">MaskHead (class in cyto_dl.nn.head.mask_head)</a>
+</li>
+        <li><a href="cyto_dl.nn.vits.blocks.masked_unit_attention.html#cyto_dl.nn.vits.blocks.masked_unit_attention.MaskUnitAttention">MaskUnitAttention (class in cyto_dl.nn.vits.blocks.masked_unit_attention)</a>
+</li>
+        <li><a href="cyto_dl.nn.vits.utils.html#cyto_dl.nn.vits.utils.match_tuple_dimensions">match_tuple_dimensions() (in module cyto_dl.nn.vits.utils)</a>
 </li>
         <li><a href="cyto_dl.image.transforms.project.html#cyto_dl.image.transforms.project.MaxProjectd">MaxProjectd (class in cyto_dl.image.transforms.project)</a>
 </li>
@@ -2375,14 +2493,24 @@ <h2>M</h2>
           <li><a href="cyto_dl.nn.vits.blocks.cross_attention.html#module-cyto_dl.nn.vits.blocks.cross_attention">cyto_dl.nn.vits.blocks.cross_attention</a>
 </li>
           <li><a href="cyto_dl.nn.vits.blocks.intermediate_weigher.html#module-cyto_dl.nn.vits.blocks.intermediate_weigher">cyto_dl.nn.vits.blocks.intermediate_weigher</a>
+</li>
+          <li><a href="cyto_dl.nn.vits.blocks.masked_unit_attention.html#module-cyto_dl.nn.vits.blocks.masked_unit_attention">cyto_dl.nn.vits.blocks.masked_unit_attention</a>
 </li>
           <li><a href="cyto_dl.nn.vits.blocks.patchify.html#module-cyto_dl.nn.vits.blocks.patchify">cyto_dl.nn.vits.blocks.patchify</a>
 </li>
-          <li><a href="cyto_dl.nn.vits.cross_mae.html#module-cyto_dl.nn.vits.cross_mae">cyto_dl.nn.vits.cross_mae</a>
+          <li><a href="cyto_dl.nn.vits.blocks.patchify.patchify.html#module-cyto_dl.nn.vits.blocks.patchify.patchify">cyto_dl.nn.vits.blocks.patchify.patchify</a>
+</li>
+          <li><a href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html#module-cyto_dl.nn.vits.blocks.patchify.patchify_base">cyto_dl.nn.vits.blocks.patchify.patchify_base</a>
+</li>
+          <li><a href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html#module-cyto_dl.nn.vits.blocks.patchify.patchify_hiera">cyto_dl.nn.vits.blocks.patchify.patchify_hiera</a>
+</li>
+          <li><a href="cyto_dl.nn.vits.decoder.html#module-cyto_dl.nn.vits.decoder">cyto_dl.nn.vits.decoder</a>
 </li>
-          <li><a href="cyto_dl.nn.vits.jepa.html#module-cyto_dl.nn.vits.jepa">cyto_dl.nn.vits.jepa</a>
+          <li><a href="cyto_dl.nn.vits.encoder.html#module-cyto_dl.nn.vits.encoder">cyto_dl.nn.vits.encoder</a>
 </li>
           <li><a href="cyto_dl.nn.vits.mae.html#module-cyto_dl.nn.vits.mae">cyto_dl.nn.vits.mae</a>
+</li>
+          <li><a href="cyto_dl.nn.vits.predictor.html#module-cyto_dl.nn.vits.predictor">cyto_dl.nn.vits.predictor</a>
 </li>
           <li><a href="cyto_dl.nn.vits.seg.html#module-cyto_dl.nn.vits.seg">cyto_dl.nn.vits.seg</a>
 </li>
@@ -2529,7 +2657,11 @@ <h2>P</h2>
 </li>
         <li><a href="cyto_dl.datamodules.dataframe.utils.html#cyto_dl.datamodules.dataframe.utils.parse_transforms">parse_transforms() (in module cyto_dl.datamodules.dataframe.utils)</a>
 </li>
-        <li><a href="cyto_dl.nn.vits.blocks.patchify.html#cyto_dl.nn.vits.blocks.patchify.Patchify">Patchify (class in cyto_dl.nn.vits.blocks.patchify)</a>
+        <li><a href="cyto_dl.nn.vits.blocks.patchify.patchify.html#cyto_dl.nn.vits.blocks.patchify.patchify.Patchify">Patchify (class in cyto_dl.nn.vits.blocks.patchify.patchify)</a>
+</li>
+        <li><a href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html#cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase">PatchifyBase (class in cyto_dl.nn.vits.blocks.patchify.patchify_base)</a>
+</li>
+        <li><a href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html#cyto_dl.nn.vits.blocks.patchify.patchify_hiera.PatchifyHiera">PatchifyHiera (class in cyto_dl.nn.vits.blocks.patchify.patchify_hiera)</a>
 </li>
         <li><a href="cyto_dl.api.data.html#cyto_dl.api.data.PatchSize">PatchSize (class in cyto_dl.api.data)</a>
 </li>
@@ -2546,11 +2678,11 @@ <h2>P</h2>
         <li><a href="cyto_dl.models.contrastive.contrastive.html#cyto_dl.models.contrastive.contrastive.Contrastive.plot_classes">plot_classes() (cyto_dl.models.contrastive.contrastive.Contrastive method)</a>
 </li>
         <li><a href="cyto_dl.models.contrastive.contrastive.html#cyto_dl.models.contrastive.contrastive.Contrastive.plot_neighbors">plot_neighbors() (cyto_dl.models.contrastive.contrastive.Contrastive method)</a>
-</li>
-        <li><a href="cyto_dl.image.io.polygon_loader.html#cyto_dl.image.io.polygon_loader.PolygonLoaderd">PolygonLoaderd (class in cyto_dl.image.io.polygon_loader)</a>
 </li>
     </ul></td>
     <td style="width: 33%; vertical-align: top;"><ul>
+        <li><a href="cyto_dl.image.io.polygon_loader.html#cyto_dl.image.io.polygon_loader.PolygonLoaderd">PolygonLoaderd (class in cyto_dl.image.io.polygon_loader)</a>
+</li>
         <li><a href="cyto_dl.nn.track_sequence_predictor.html#cyto_dl.nn.track_sequence_predictor.positionalencoding1d">positionalencoding1d() (in module cyto_dl.nn.track_sequence_predictor)</a>
 </li>
         <li><a href="cyto_dl.callbacks.csv_saver.html#cyto_dl.callbacks.csv_saver.CSVSaver.pred_to_csv">pred_to_csv() (cyto_dl.callbacks.csv_saver.CSVSaver method)</a>
@@ -2591,7 +2723,7 @@ <h2>P</h2>
           <li><a href="cyto_dl.models.jepa.jepa_base.html#cyto_dl.models.jepa.jepa_base.JEPABase.predict_step">(cyto_dl.models.jepa.jepa_base.JEPABase method)</a>
 </li>
         </ul></li>
-        <li><a href="cyto_dl.nn.vits.jepa.html#cyto_dl.nn.vits.jepa.JEPAPredictor.predict_target_features">predict_target_features() (cyto_dl.nn.vits.jepa.JEPAPredictor method)</a>
+        <li><a href="cyto_dl.nn.vits.predictor.html#cyto_dl.nn.vits.predictor.JEPAPredictor.predict_target_features">predict_target_features() (cyto_dl.nn.vits.predictor.JEPAPredictor method)</a>
 </li>
         <li><a href="cyto_dl.datamodules.smartcache.html#cyto_dl.datamodules.smartcache.SmartcacheDatamodule.prepare_data">prepare_data() (cyto_dl.datamodules.smartcache.SmartcacheDatamodule method)</a>
 </li>
@@ -2619,7 +2751,7 @@ <h2>Q</h2>
   <h2>R</h2>
   <table style="width: 100%" class="indextable genindextable"><tr>
     <td style="width: 33%; vertical-align: top;"><ul>
-        <li><a href="cyto_dl.nn.vits.blocks.patchify.html#cyto_dl.nn.vits.blocks.patchify.random_indexes">random_indexes() (in module cyto_dl.nn.vits.blocks.patchify)</a>
+        <li><a href="cyto_dl.nn.vits.utils.html#cyto_dl.nn.vits.utils.random_indexes">random_indexes() (in module cyto_dl.nn.vits.utils)</a>
 </li>
         <li><a href="cyto_dl.image.transforms.multiscale_cropper.html#cyto_dl.image.transforms.multiscale_cropper.RandomMultiScaleCropd">RandomMultiScaleCropd (class in cyto_dl.image.transforms.multiscale_cropper)</a>
 </li>
@@ -2737,10 +2869,10 @@ <h2>S</h2>
 </li>
         <li><a href="cyto_dl.api.cyto_dl_model.cyto_dl_base_model.html#cyto_dl.api.cyto_dl_model.cyto_dl_base_model.CytoDLBaseModel.set_experiment_name">set_experiment_name() (cyto_dl.api.cyto_dl_model.cyto_dl_base_model.CytoDLBaseModel method)</a>
 </li>
-    </ul></td>
-    <td style="width: 33%; vertical-align: top;"><ul>
         <li><a href="cyto_dl.api.cyto_dl_model.segmentation_plugin_model.html#cyto_dl.api.cyto_dl_model.segmentation_plugin_model.SegmentationPluginModel.set_hardware_type">set_hardware_type() (cyto_dl.api.cyto_dl_model.segmentation_plugin_model.SegmentationPluginModel method)</a>
 </li>
+    </ul></td>
+    <td style="width: 33%; vertical-align: top;"><ul>
         <li><a href="cyto_dl.api.cyto_dl_model.segmentation_plugin_model.html#cyto_dl.api.cyto_dl_model.segmentation_plugin_model.SegmentationPluginModel.set_input_channel">set_input_channel() (cyto_dl.api.cyto_dl_model.segmentation_plugin_model.SegmentationPluginModel method)</a>
 </li>
         <li><a href="cyto_dl.api.cyto_dl_model.segmentation_plugin_model.html#cyto_dl.api.cyto_dl_model.segmentation_plugin_model.SegmentationPluginModel.set_manifest_column_names">set_manifest_column_names() (cyto_dl.api.cyto_dl_model.segmentation_plugin_model.SegmentationPluginModel method)</a>
@@ -2772,6 +2904,8 @@ <h2>S</h2>
         <li><a href="cyto_dl.datamodules.smartcache.html#cyto_dl.datamodules.smartcache.SmartcacheDatamodule">SmartcacheDatamodule (class in cyto_dl.datamodules.smartcache)</a>
 </li>
         <li><a href="cyto_dl.models.im2im.utils.instance_seg.html#cyto_dl.models.im2im.utils.instance_seg.InstanceSegPreprocessd.smooth_embedding">smooth_embedding() (cyto_dl.models.im2im.utils.instance_seg.InstanceSegPreprocessd method)</a>
+</li>
+        <li><a href="cyto_dl.nn.vits.encoder.html#cyto_dl.nn.vits.encoder.SpatialMerger">SpatialMerger (class in cyto_dl.nn.vits.encoder)</a>
 </li>
         <li><a href="cyto_dl.image.io.aicsimage_loader.html#cyto_dl.image.io.aicsimage_loader.AICSImageLoaderd.split_args">split_args() (cyto_dl.image.io.aicsimage_loader.AICSImageLoaderd method)</a>
 </li>
@@ -2800,6 +2934,8 @@ <h2>T</h2>
   <table style="width: 100%" class="indextable genindextable"><tr>
     <td style="width: 33%; vertical-align: top;"><ul>
         <li><a href="cyto_dl.nn.vits.utils.html#cyto_dl.nn.vits.utils.take_indexes">take_indexes() (in module cyto_dl.nn.vits.utils)</a>
+</li>
+        <li><a href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html#cyto_dl.nn.vits.blocks.patchify.patchify_hiera.take_indexes_mask">take_indexes_mask() (in module cyto_dl.nn.vits.blocks.patchify.patchify_hiera)</a>
 </li>
         <li><a href="cyto_dl.utils.template_utils.html#cyto_dl.utils.template_utils.task_wrapper">task_wrapper() (in module cyto_dl.utils.template_utils)</a>
 </li>
diff --git a/index.html b/index.html
index 9c2473782..0446686f6 100644
--- a/index.html
+++ b/index.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/installation.html b/installation.html
index 582845b87..5d84900c4 100644
--- a/installation.html
+++ b/installation.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/modules.html b/modules.html
index 84d334923..7d23dcaff 100644
--- a/modules.html
+++ b/modules.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/objects.inv b/objects.inv
index 1669e4d12..7c06b0d7f 100644
Binary files a/objects.inv and b/objects.inv differ
diff --git a/overview.html b/overview.html
index e771e5284..d7ed3bd74 100644
--- a/overview.html
+++ b/overview.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/py-modindex.html b/py-modindex.html
index 2efd7a82f..f401b403c 100644
--- a/py-modindex.html
+++ b/py-modindex.html
@@ -348,14 +348,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -367,15 +374,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
@@ -1152,6 +1159,12 @@ <h1>Python Module Index</h1>
         <a href="cyto_dl.nn.vits.blocks.intermediate_weigher.html#module-cyto_dl.nn.vits.blocks.intermediate_weigher"><code class="xref">cyto_dl.nn.vits.blocks.intermediate_weigher</code></a></td><td>
     <em></em></td>
   </tr>
+  <tr class="cg-1">
+    <td></td>
+    <td>&#160;&#160;&#160;
+        <a href="cyto_dl.nn.vits.blocks.masked_unit_attention.html#module-cyto_dl.nn.vits.blocks.masked_unit_attention"><code class="xref">cyto_dl.nn.vits.blocks.masked_unit_attention</code></a></td><td>
+    <em></em></td>
+  </tr>
   <tr class="cg-1">
     <td></td>
     <td>&#160;&#160;&#160;
@@ -1161,13 +1174,31 @@ <h1>Python Module Index</h1>
   <tr class="cg-1">
     <td></td>
     <td>&#160;&#160;&#160;
-        <a href="cyto_dl.nn.vits.cross_mae.html#module-cyto_dl.nn.vits.cross_mae"><code class="xref">cyto_dl.nn.vits.cross_mae</code></a></td><td>
+        <a href="cyto_dl.nn.vits.blocks.patchify.patchify.html#module-cyto_dl.nn.vits.blocks.patchify.patchify"><code class="xref">cyto_dl.nn.vits.blocks.patchify.patchify</code></a></td><td>
     <em></em></td>
   </tr>
   <tr class="cg-1">
     <td></td>
     <td>&#160;&#160;&#160;
-        <a href="cyto_dl.nn.vits.jepa.html#module-cyto_dl.nn.vits.jepa"><code class="xref">cyto_dl.nn.vits.jepa</code></a></td><td>
+        <a href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html#module-cyto_dl.nn.vits.blocks.patchify.patchify_base"><code class="xref">cyto_dl.nn.vits.blocks.patchify.patchify_base</code></a></td><td>
+    <em></em></td>
+  </tr>
+  <tr class="cg-1">
+    <td></td>
+    <td>&#160;&#160;&#160;
+        <a href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html#module-cyto_dl.nn.vits.blocks.patchify.patchify_hiera"><code class="xref">cyto_dl.nn.vits.blocks.patchify.patchify_hiera</code></a></td><td>
+    <em></em></td>
+  </tr>
+  <tr class="cg-1">
+    <td></td>
+    <td>&#160;&#160;&#160;
+        <a href="cyto_dl.nn.vits.decoder.html#module-cyto_dl.nn.vits.decoder"><code class="xref">cyto_dl.nn.vits.decoder</code></a></td><td>
+    <em></em></td>
+  </tr>
+  <tr class="cg-1">
+    <td></td>
+    <td>&#160;&#160;&#160;
+        <a href="cyto_dl.nn.vits.encoder.html#module-cyto_dl.nn.vits.encoder"><code class="xref">cyto_dl.nn.vits.encoder</code></a></td><td>
     <em></em></td>
   </tr>
   <tr class="cg-1">
@@ -1176,6 +1207,12 @@ <h1>Python Module Index</h1>
         <a href="cyto_dl.nn.vits.mae.html#module-cyto_dl.nn.vits.mae"><code class="xref">cyto_dl.nn.vits.mae</code></a></td><td>
     <em></em></td>
   </tr>
+  <tr class="cg-1">
+    <td></td>
+    <td>&#160;&#160;&#160;
+        <a href="cyto_dl.nn.vits.predictor.html#module-cyto_dl.nn.vits.predictor"><code class="xref">cyto_dl.nn.vits.predictor</code></a></td><td>
+    <em></em></td>
+  </tr>
   <tr class="cg-1">
     <td></td>
     <td>&#160;&#160;&#160;
diff --git a/search.html b/search.html
index 7bd3abcc2..2e4bf3348 100644
--- a/search.html
+++ b/search.html
@@ -347,14 +347,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -366,15 +373,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>
diff --git a/searchindex.js b/searchindex.js
index ad09c50eb..46a04aa56 100644
--- a/searchindex.js
+++ b/searchindex.js
@@ -1 +1 @@
-Search.setIndex({"docnames": ["api", "cyto_dl", "cyto_dl.api", "cyto_dl.api.cyto_dl_model", "cyto_dl.api.cyto_dl_model.cyto_dl_base_model", "cyto_dl.api.cyto_dl_model.segmentation_plugin_model", "cyto_dl.api.data", "cyto_dl.api.model", "cyto_dl.callbacks", "cyto_dl.callbacks.callback_utils", "cyto_dl.callbacks.csv_saver", "cyto_dl.callbacks.image_saver", "cyto_dl.callbacks.latent_walk", "cyto_dl.callbacks.layer_freeze", "cyto_dl.callbacks.model_utils", "cyto_dl.callbacks.outlier_detection", "cyto_dl.compile", "cyto_dl.dataframe", "cyto_dl.dataframe.readers", "cyto_dl.dataframe.transforms", "cyto_dl.dataframe.transforms.filter", "cyto_dl.dataframe.transforms.group_cols", "cyto_dl.dataframe.transforms.misc", "cyto_dl.dataframe.transforms.split", "cyto_dl.datamodules", "cyto_dl.datamodules.array", "cyto_dl.datamodules.data_dict", "cyto_dl.datamodules.dataframe", "cyto_dl.datamodules.dataframe.dataframe_datamodule", "cyto_dl.datamodules.dataframe.grouped_dataframe_datamodule", "cyto_dl.datamodules.dataframe.utils", "cyto_dl.datamodules.folder", "cyto_dl.datamodules.multidim_image", "cyto_dl.datamodules.smartcache", "cyto_dl.datamodules.torchvision", "cyto_dl.eval", "cyto_dl.image", "cyto_dl.image.io", "cyto_dl.image.io.aicsimage_loader", "cyto_dl.image.io.monai_bio_reader", "cyto_dl.image.io.numpy_reader", "cyto_dl.image.io.ome_zarr_reader", "cyto_dl.image.io.polygon_loader", "cyto_dl.image.io.skimage_reader", "cyto_dl.image.transforms", "cyto_dl.image.transforms.add_meta", "cyto_dl.image.transforms.bright_sampler", "cyto_dl.image.transforms.clip", "cyto_dl.image.transforms.contrastadjust", "cyto_dl.image.transforms.generate_jepa_masks", "cyto_dl.image.transforms.merge", "cyto_dl.image.transforms.multiscale_cropper", "cyto_dl.image.transforms.pad", "cyto_dl.image.transforms.project", "cyto_dl.image.transforms.rotation_mask_transform", "cyto_dl.image.transforms.save", "cyto_dl.image.transforms.so2_random_rotation", "cyto_dl.image.transforms.track_transforms", "cyto_dl.loggers", "cyto_dl.loggers.mlflow", "cyto_dl.models", "cyto_dl.models.base_model", "cyto_dl.models.basic_model", "cyto_dl.models.classification", "cyto_dl.models.classification.classification", "cyto_dl.models.classification.timepoint_classification", "cyto_dl.models.contrastive", "cyto_dl.models.contrastive.contrastive", "cyto_dl.models.handlers", "cyto_dl.models.handlers.base_handler", "cyto_dl.models.handlers.image_handler", "cyto_dl.models.handlers.load_image_patch", "cyto_dl.models.im2im", "cyto_dl.models.im2im.gan", "cyto_dl.models.im2im.multi_task", "cyto_dl.models.im2im.utils", "cyto_dl.models.im2im.utils.instance_seg", "cyto_dl.models.im2im.utils.noise_annealer", "cyto_dl.models.im2im.utils.postprocessing", "cyto_dl.models.im2im.utils.postprocessing.act_thresh_label", "cyto_dl.models.im2im.utils.postprocessing.auto_thresh", "cyto_dl.models.im2im.utils.postprocessing.dict_to_im", "cyto_dl.models.jepa", "cyto_dl.models.jepa.ijepa", "cyto_dl.models.jepa.iwm", "cyto_dl.models.jepa.jepa_base", "cyto_dl.models.utils", "cyto_dl.models.utils.mlflow", "cyto_dl.models.utils.utils", "cyto_dl.models.vae", "cyto_dl.models.vae.base_vae", "cyto_dl.models.vae.image_canon_vae", "cyto_dl.models.vae.image_encoder", "cyto_dl.models.vae.image_vae", "cyto_dl.models.vae.latent_loss_vae", "cyto_dl.models.vae.o2_spharm_vae", "cyto_dl.models.vae.o2_spharm_vae.o2_spharm_encoder", "cyto_dl.models.vae.o2_spharm_vae.o2_spharm_vae", "cyto_dl.models.vae.point_cloud_vae", "cyto_dl.models.vae.priors", "cyto_dl.models.vae.priors.abstract_prior", "cyto_dl.models.vae.priors.gaussian", "cyto_dl.models.vae.priors.identity_prior", "cyto_dl.models.vae.priors.joint_prior", "cyto_dl.models.vae.tabular_vae", "cyto_dl.nn", "cyto_dl.nn.discriminators", "cyto_dl.nn.discriminators.multi_scale_discriminator", "cyto_dl.nn.discriminators.n_layer_discriminator", "cyto_dl.nn.head", "cyto_dl.nn.head.base_head", "cyto_dl.nn.head.gan_head", "cyto_dl.nn.head.gan_head_superres", "cyto_dl.nn.head.mae_head", "cyto_dl.nn.head.mask_head", "cyto_dl.nn.head.res_blocks_head", "cyto_dl.nn.head.vic_reg", "cyto_dl.nn.hr_skip", "cyto_dl.nn.losses", "cyto_dl.nn.losses.adversarial_loss", "cyto_dl.nn.losses.chamfer_loss", "cyto_dl.nn.losses.continuous_bernoulli", "cyto_dl.nn.losses.cosine_loss", "cyto_dl.nn.losses.gan_loss", "cyto_dl.nn.losses.gaussian_nll_loss", "cyto_dl.nn.losses.geomloss", "cyto_dl.nn.losses.loss_wrapper", "cyto_dl.nn.losses.spharm_loss", "cyto_dl.nn.losses.threshold_loss", "cyto_dl.nn.losses.vic_reg", "cyto_dl.nn.losses.weibull", "cyto_dl.nn.losses.weighted_mse_loss", "cyto_dl.nn.mlp", "cyto_dl.nn.point_cloud", "cyto_dl.nn.point_cloud.dgcnn", "cyto_dl.nn.point_cloud.folding_net", "cyto_dl.nn.point_cloud.graph_functions", "cyto_dl.nn.point_cloud.vnn", "cyto_dl.nn.res_unit", "cyto_dl.nn.spatial_transformer", "cyto_dl.nn.track_sequence_predictor", "cyto_dl.nn.vits", "cyto_dl.nn.vits.blocks", "cyto_dl.nn.vits.blocks.cross_attention", "cyto_dl.nn.vits.blocks.intermediate_weigher", "cyto_dl.nn.vits.blocks.patchify", "cyto_dl.nn.vits.cross_mae", "cyto_dl.nn.vits.jepa", "cyto_dl.nn.vits.mae", "cyto_dl.nn.vits.seg", "cyto_dl.nn.vits.utils", "cyto_dl.point_cloud", "cyto_dl.point_cloud.io", "cyto_dl.point_cloud.io.read_pcloud", "cyto_dl.train", "cyto_dl.utils", "cyto_dl.utils.array", "cyto_dl.utils.checkpoint", "cyto_dl.utils.config", "cyto_dl.utils.download_test_data", "cyto_dl.utils.dummy_dataset", "cyto_dl.utils.pylogger", "cyto_dl.utils.rich_utils", "cyto_dl.utils.rotation", "cyto_dl.utils.spharm", "cyto_dl.utils.spharm.reconstruction", "cyto_dl.utils.spharm.rotation", "cyto_dl.utils.template_utils", "index", "installation", "modules", "overview", "using_examples"], "filenames": ["api.rst", "cyto_dl.rst", "cyto_dl.api.rst", "cyto_dl.api.cyto_dl_model.rst", "cyto_dl.api.cyto_dl_model.cyto_dl_base_model.rst", "cyto_dl.api.cyto_dl_model.segmentation_plugin_model.rst", "cyto_dl.api.data.rst", "cyto_dl.api.model.rst", "cyto_dl.callbacks.rst", "cyto_dl.callbacks.callback_utils.rst", "cyto_dl.callbacks.csv_saver.rst", "cyto_dl.callbacks.image_saver.rst", "cyto_dl.callbacks.latent_walk.rst", "cyto_dl.callbacks.layer_freeze.rst", "cyto_dl.callbacks.model_utils.rst", "cyto_dl.callbacks.outlier_detection.rst", "cyto_dl.compile.rst", "cyto_dl.dataframe.rst", "cyto_dl.dataframe.readers.rst", "cyto_dl.dataframe.transforms.rst", "cyto_dl.dataframe.transforms.filter.rst", "cyto_dl.dataframe.transforms.group_cols.rst", "cyto_dl.dataframe.transforms.misc.rst", "cyto_dl.dataframe.transforms.split.rst", "cyto_dl.datamodules.rst", "cyto_dl.datamodules.array.rst", "cyto_dl.datamodules.data_dict.rst", "cyto_dl.datamodules.dataframe.rst", "cyto_dl.datamodules.dataframe.dataframe_datamodule.rst", "cyto_dl.datamodules.dataframe.grouped_dataframe_datamodule.rst", "cyto_dl.datamodules.dataframe.utils.rst", "cyto_dl.datamodules.folder.rst", "cyto_dl.datamodules.multidim_image.rst", "cyto_dl.datamodules.smartcache.rst", "cyto_dl.datamodules.torchvision.rst", "cyto_dl.eval.rst", "cyto_dl.image.rst", "cyto_dl.image.io.rst", "cyto_dl.image.io.aicsimage_loader.rst", "cyto_dl.image.io.monai_bio_reader.rst", "cyto_dl.image.io.numpy_reader.rst", "cyto_dl.image.io.ome_zarr_reader.rst", "cyto_dl.image.io.polygon_loader.rst", "cyto_dl.image.io.skimage_reader.rst", "cyto_dl.image.transforms.rst", "cyto_dl.image.transforms.add_meta.rst", "cyto_dl.image.transforms.bright_sampler.rst", "cyto_dl.image.transforms.clip.rst", "cyto_dl.image.transforms.contrastadjust.rst", "cyto_dl.image.transforms.generate_jepa_masks.rst", "cyto_dl.image.transforms.merge.rst", "cyto_dl.image.transforms.multiscale_cropper.rst", "cyto_dl.image.transforms.pad.rst", "cyto_dl.image.transforms.project.rst", "cyto_dl.image.transforms.rotation_mask_transform.rst", "cyto_dl.image.transforms.save.rst", "cyto_dl.image.transforms.so2_random_rotation.rst", "cyto_dl.image.transforms.track_transforms.rst", "cyto_dl.loggers.rst", "cyto_dl.loggers.mlflow.rst", "cyto_dl.models.rst", "cyto_dl.models.base_model.rst", "cyto_dl.models.basic_model.rst", "cyto_dl.models.classification.rst", "cyto_dl.models.classification.classification.rst", "cyto_dl.models.classification.timepoint_classification.rst", "cyto_dl.models.contrastive.rst", "cyto_dl.models.contrastive.contrastive.rst", "cyto_dl.models.handlers.rst", "cyto_dl.models.handlers.base_handler.rst", "cyto_dl.models.handlers.image_handler.rst", "cyto_dl.models.handlers.load_image_patch.rst", "cyto_dl.models.im2im.rst", "cyto_dl.models.im2im.gan.rst", "cyto_dl.models.im2im.multi_task.rst", "cyto_dl.models.im2im.utils.rst", "cyto_dl.models.im2im.utils.instance_seg.rst", "cyto_dl.models.im2im.utils.noise_annealer.rst", "cyto_dl.models.im2im.utils.postprocessing.rst", "cyto_dl.models.im2im.utils.postprocessing.act_thresh_label.rst", "cyto_dl.models.im2im.utils.postprocessing.auto_thresh.rst", "cyto_dl.models.im2im.utils.postprocessing.dict_to_im.rst", "cyto_dl.models.jepa.rst", "cyto_dl.models.jepa.ijepa.rst", "cyto_dl.models.jepa.iwm.rst", "cyto_dl.models.jepa.jepa_base.rst", "cyto_dl.models.utils.rst", "cyto_dl.models.utils.mlflow.rst", "cyto_dl.models.utils.utils.rst", "cyto_dl.models.vae.rst", "cyto_dl.models.vae.base_vae.rst", "cyto_dl.models.vae.image_canon_vae.rst", "cyto_dl.models.vae.image_encoder.rst", "cyto_dl.models.vae.image_vae.rst", "cyto_dl.models.vae.latent_loss_vae.rst", "cyto_dl.models.vae.o2_spharm_vae.rst", "cyto_dl.models.vae.o2_spharm_vae.o2_spharm_encoder.rst", "cyto_dl.models.vae.o2_spharm_vae.o2_spharm_vae.rst", "cyto_dl.models.vae.point_cloud_vae.rst", "cyto_dl.models.vae.priors.rst", "cyto_dl.models.vae.priors.abstract_prior.rst", "cyto_dl.models.vae.priors.gaussian.rst", "cyto_dl.models.vae.priors.identity_prior.rst", "cyto_dl.models.vae.priors.joint_prior.rst", "cyto_dl.models.vae.tabular_vae.rst", "cyto_dl.nn.rst", "cyto_dl.nn.discriminators.rst", "cyto_dl.nn.discriminators.multi_scale_discriminator.rst", "cyto_dl.nn.discriminators.n_layer_discriminator.rst", "cyto_dl.nn.head.rst", "cyto_dl.nn.head.base_head.rst", "cyto_dl.nn.head.gan_head.rst", "cyto_dl.nn.head.gan_head_superres.rst", "cyto_dl.nn.head.mae_head.rst", "cyto_dl.nn.head.mask_head.rst", "cyto_dl.nn.head.res_blocks_head.rst", "cyto_dl.nn.head.vic_reg.rst", "cyto_dl.nn.hr_skip.rst", "cyto_dl.nn.losses.rst", "cyto_dl.nn.losses.adversarial_loss.rst", "cyto_dl.nn.losses.chamfer_loss.rst", "cyto_dl.nn.losses.continuous_bernoulli.rst", "cyto_dl.nn.losses.cosine_loss.rst", "cyto_dl.nn.losses.gan_loss.rst", "cyto_dl.nn.losses.gaussian_nll_loss.rst", "cyto_dl.nn.losses.geomloss.rst", "cyto_dl.nn.losses.loss_wrapper.rst", "cyto_dl.nn.losses.spharm_loss.rst", "cyto_dl.nn.losses.threshold_loss.rst", "cyto_dl.nn.losses.vic_reg.rst", "cyto_dl.nn.losses.weibull.rst", "cyto_dl.nn.losses.weighted_mse_loss.rst", "cyto_dl.nn.mlp.rst", "cyto_dl.nn.point_cloud.rst", "cyto_dl.nn.point_cloud.dgcnn.rst", "cyto_dl.nn.point_cloud.folding_net.rst", "cyto_dl.nn.point_cloud.graph_functions.rst", "cyto_dl.nn.point_cloud.vnn.rst", "cyto_dl.nn.res_unit.rst", "cyto_dl.nn.spatial_transformer.rst", "cyto_dl.nn.track_sequence_predictor.rst", "cyto_dl.nn.vits.rst", "cyto_dl.nn.vits.blocks.rst", "cyto_dl.nn.vits.blocks.cross_attention.rst", "cyto_dl.nn.vits.blocks.intermediate_weigher.rst", "cyto_dl.nn.vits.blocks.patchify.rst", "cyto_dl.nn.vits.cross_mae.rst", "cyto_dl.nn.vits.jepa.rst", "cyto_dl.nn.vits.mae.rst", "cyto_dl.nn.vits.seg.rst", "cyto_dl.nn.vits.utils.rst", "cyto_dl.point_cloud.rst", "cyto_dl.point_cloud.io.rst", "cyto_dl.point_cloud.io.read_pcloud.rst", "cyto_dl.train.rst", "cyto_dl.utils.rst", "cyto_dl.utils.array.rst", "cyto_dl.utils.checkpoint.rst", "cyto_dl.utils.config.rst", "cyto_dl.utils.download_test_data.rst", "cyto_dl.utils.dummy_dataset.rst", "cyto_dl.utils.pylogger.rst", "cyto_dl.utils.rich_utils.rst", "cyto_dl.utils.rotation.rst", "cyto_dl.utils.spharm.rst", "cyto_dl.utils.spharm.reconstruction.rst", "cyto_dl.utils.spharm.rotation.rst", "cyto_dl.utils.template_utils.rst", "index.rst", "installation.rst", "modules.rst", "overview.md", "using_examples.rst"], "titles": ["Using the API", "cyto_dl package", "cyto_dl.api package", "cyto_dl.api.cyto_dl_model package", "cyto_dl.api.cyto_dl_model.cyto_dl_base_model module", "cyto_dl.api.cyto_dl_model.segmentation_plugin_model module", "cyto_dl.api.data module", "cyto_dl.api.model module", "cyto_dl.callbacks package", "cyto_dl.callbacks.callback_utils module", "cyto_dl.callbacks.csv_saver module", "cyto_dl.callbacks.image_saver module", "cyto_dl.callbacks.latent_walk module", "cyto_dl.callbacks.layer_freeze module", "cyto_dl.callbacks.model_utils module", "cyto_dl.callbacks.outlier_detection module", "cyto_dl.compile module", "cyto_dl.dataframe package", "cyto_dl.dataframe.readers module", "cyto_dl.dataframe.transforms package", "cyto_dl.dataframe.transforms.filter module", "cyto_dl.dataframe.transforms.group_cols module", "cyto_dl.dataframe.transforms.misc module", "cyto_dl.dataframe.transforms.split module", "cyto_dl.datamodules package", "cyto_dl.datamodules.array module", "cyto_dl.datamodules.data_dict module", "cyto_dl.datamodules.dataframe package", "cyto_dl.datamodules.dataframe.dataframe_datamodule module", "cyto_dl.datamodules.dataframe.grouped_dataframe_datamodule module", "cyto_dl.datamodules.dataframe.utils module", "cyto_dl.datamodules.folder module", "cyto_dl.datamodules.multidim_image module", "cyto_dl.datamodules.smartcache module", "cyto_dl.datamodules.torchvision module", "cyto_dl.eval module", "cyto_dl.image package", "cyto_dl.image.io package", "cyto_dl.image.io.aicsimage_loader module", "cyto_dl.image.io.monai_bio_reader module", "cyto_dl.image.io.numpy_reader module", "cyto_dl.image.io.ome_zarr_reader module", "cyto_dl.image.io.polygon_loader module", "cyto_dl.image.io.skimage_reader module", "cyto_dl.image.transforms package", "cyto_dl.image.transforms.add_meta module", "cyto_dl.image.transforms.bright_sampler module", "cyto_dl.image.transforms.clip module", "cyto_dl.image.transforms.contrastadjust module", "cyto_dl.image.transforms.generate_jepa_masks module", "cyto_dl.image.transforms.merge module", "cyto_dl.image.transforms.multiscale_cropper module", "cyto_dl.image.transforms.pad module", "cyto_dl.image.transforms.project module", "cyto_dl.image.transforms.rotation_mask_transform module", "cyto_dl.image.transforms.save module", "cyto_dl.image.transforms.so2_random_rotation module", "cyto_dl.image.transforms.track_transforms module", "cyto_dl.loggers package", "cyto_dl.loggers.mlflow module", "cyto_dl.models package", "cyto_dl.models.base_model module", "cyto_dl.models.basic_model module", "cyto_dl.models.classification package", "cyto_dl.models.classification.classification module", "cyto_dl.models.classification.timepoint_classification module", "cyto_dl.models.contrastive package", "cyto_dl.models.contrastive.contrastive module", "cyto_dl.models.handlers package", "cyto_dl.models.handlers.base_handler module", "cyto_dl.models.handlers.image_handler module", "cyto_dl.models.handlers.load_image_patch module", "cyto_dl.models.im2im package", "cyto_dl.models.im2im.gan module", "cyto_dl.models.im2im.multi_task module", "cyto_dl.models.im2im.utils package", "cyto_dl.models.im2im.utils.instance_seg module", "cyto_dl.models.im2im.utils.noise_annealer module", "cyto_dl.models.im2im.utils.postprocessing package", "cyto_dl.models.im2im.utils.postprocessing.act_thresh_label module", "cyto_dl.models.im2im.utils.postprocessing.auto_thresh module", "cyto_dl.models.im2im.utils.postprocessing.dict_to_im module", "cyto_dl.models.jepa package", "cyto_dl.models.jepa.ijepa module", "cyto_dl.models.jepa.iwm module", "cyto_dl.models.jepa.jepa_base module", "cyto_dl.models.utils package", "cyto_dl.models.utils.mlflow module", "cyto_dl.models.utils.utils module", "cyto_dl.models.vae package", "cyto_dl.models.vae.base_vae module", "cyto_dl.models.vae.image_canon_vae module", "cyto_dl.models.vae.image_encoder module", "cyto_dl.models.vae.image_vae module", "cyto_dl.models.vae.latent_loss_vae module", "cyto_dl.models.vae.o2_spharm_vae package", "cyto_dl.models.vae.o2_spharm_vae.o2_spharm_encoder module", "cyto_dl.models.vae.o2_spharm_vae.o2_spharm_vae module", "cyto_dl.models.vae.point_cloud_vae module", "cyto_dl.models.vae.priors package", "cyto_dl.models.vae.priors.abstract_prior module", "cyto_dl.models.vae.priors.gaussian module", "cyto_dl.models.vae.priors.identity_prior module", "cyto_dl.models.vae.priors.joint_prior module", "cyto_dl.models.vae.tabular_vae module", "cyto_dl.nn package", "cyto_dl.nn.discriminators package", "cyto_dl.nn.discriminators.multi_scale_discriminator module", "cyto_dl.nn.discriminators.n_layer_discriminator module", "cyto_dl.nn.head package", "cyto_dl.nn.head.base_head module", "cyto_dl.nn.head.gan_head module", "cyto_dl.nn.head.gan_head_superres module", "cyto_dl.nn.head.mae_head module", "cyto_dl.nn.head.mask_head module", "cyto_dl.nn.head.res_blocks_head module", "cyto_dl.nn.head.vic_reg module", "cyto_dl.nn.hr_skip module", "cyto_dl.nn.losses package", "cyto_dl.nn.losses.adversarial_loss module", "cyto_dl.nn.losses.chamfer_loss module", "cyto_dl.nn.losses.continuous_bernoulli module", "cyto_dl.nn.losses.cosine_loss module", "cyto_dl.nn.losses.gan_loss module", "cyto_dl.nn.losses.gaussian_nll_loss module", "cyto_dl.nn.losses.geomloss module", "cyto_dl.nn.losses.loss_wrapper module", "cyto_dl.nn.losses.spharm_loss module", "cyto_dl.nn.losses.threshold_loss module", "cyto_dl.nn.losses.vic_reg module", "cyto_dl.nn.losses.weibull module", "cyto_dl.nn.losses.weighted_mse_loss module", "cyto_dl.nn.mlp module", "cyto_dl.nn.point_cloud package", "cyto_dl.nn.point_cloud.dgcnn module", "cyto_dl.nn.point_cloud.folding_net module", "cyto_dl.nn.point_cloud.graph_functions module", "cyto_dl.nn.point_cloud.vnn module", "cyto_dl.nn.res_unit module", "cyto_dl.nn.spatial_transformer module", "cyto_dl.nn.track_sequence_predictor module", "cyto_dl.nn.vits package", "cyto_dl.nn.vits.blocks package", "cyto_dl.nn.vits.blocks.cross_attention module", "cyto_dl.nn.vits.blocks.intermediate_weigher module", "cyto_dl.nn.vits.blocks.patchify module", "cyto_dl.nn.vits.cross_mae module", "cyto_dl.nn.vits.jepa module", "cyto_dl.nn.vits.mae module", "cyto_dl.nn.vits.seg module", "cyto_dl.nn.vits.utils module", "cyto_dl.point_cloud package", "cyto_dl.point_cloud.io package", "cyto_dl.point_cloud.io.read_pcloud module", "cyto_dl.train module", "cyto_dl.utils package", "cyto_dl.utils.array module", "cyto_dl.utils.checkpoint module", "cyto_dl.utils.config module", "cyto_dl.utils.download_test_data module", "cyto_dl.utils.dummy_dataset module", "cyto_dl.utils.pylogger module", "cyto_dl.utils.rich_utils module", "cyto_dl.utils.rotation module", "cyto_dl.utils.spharm package", "cyto_dl.utils.spharm.reconstruction module", "cyto_dl.utils.spharm.rotation module", "cyto_dl.utils.template_utils module", "Overview", "Installation", "cyto_dl", "Overview", "Modifying the Example Configs"], "terms": {"cytodl": [0, 5, 168, 171], "can": [0, 18, 20, 21, 28, 29, 31, 38, 40, 51, 52, 77, 90, 147, 150, 168, 171, 172], "import": [0, 138, 168, 171, 172], "python": [0, 21, 161, 167, 169], "packag": [0, 170], "programmat": 0, "cyto_dl": [0, 168, 171, 172], "cytodlmodel": [0, 1, 2, 7], "class": [0, 4, 5, 6, 7, 10, 11, 13, 14, 15, 21, 22, 23, 25, 26, 28, 29, 30, 31, 32, 33, 34, 38, 39, 40, 41, 42, 43, 45, 46, 47, 48, 49, 50, 51, 52, 53, 55, 57, 59, 61, 62, 64, 65, 67, 73, 74, 76, 77, 79, 80, 81, 83, 84, 85, 90, 100, 101, 102, 103, 107, 108, 110, 111, 112, 113, 114, 115, 116, 117, 119, 120, 121, 122, 123, 124, 126, 128, 129, 130, 131, 132, 138, 139, 140, 143, 144, 145, 146, 147, 148, 149, 160, 163], "addit": [0, 25, 26, 28, 29, 31, 33, 39, 61, 64, 67, 73, 74, 83, 84, 85, 90, 168, 171, 172], "command": [0, 161, 162, 167], "line": [0, 161, 162, 167, 172], "interfac": [0, 90], "we": [0, 168, 171, 172], "provid": [0, 7, 22, 32, 47, 48, 76, 126, 162, 168, 171, 172], "configur": [0, 4, 5, 7, 162, 172], "file": [0, 4, 7, 18, 28, 29, 31, 32, 33, 40, 110, 167, 168, 171, 172], "model": [0, 1, 2, 4, 51, 115, 140, 156, 157, 162, 167, 170, 172], "make": [0, 14, 77, 167, 172], "from": [0, 4, 5, 7, 14, 18, 25, 26, 28, 29, 30, 31, 32, 38, 39, 43, 45, 46, 49, 51, 57, 62, 77, 81, 112, 115, 120, 123, 138, 143, 147, 156, 162, 163, 167, 172], "dictionari": [0, 7, 21, 25, 26, 28, 29, 30, 31, 32, 38, 40, 45, 47, 48, 51, 61, 67, 76, 81, 90, 156, 172], "cfg": [0, 4, 5, 35, 154, 158, 162, 167], "data": [0, 1, 2, 7, 22, 23, 25, 26, 28, 29, 30, 31, 33, 38, 39, 41, 43, 45, 79, 83, 84, 85, 115, 156, 162, 172], "load_config_from_dict": [0, 2, 7], "yaml": [0, 4], "load_config_from_yaml": 0, "path": [0, 4, 7, 11, 18, 26, 28, 29, 31, 32, 33, 38, 40, 42, 55, 62, 87, 149, 162, 167, 168, 171, 172], "name": [0, 4, 13, 18, 21, 28, 29, 30, 31, 46, 47, 48, 50, 51, 52, 64, 147, 161, 168, 171, 172], "one": [0, 22, 28, 29, 30, 61, 167, 172], "our": [0, 172], "default": [0, 4, 5, 23, 28, 29, 41, 74, 79, 83, 84, 85, 138, 149, 168, 171, 172], "avail": [0, 29, 30], "option": [0, 13, 18, 20, 21, 23, 26, 28, 29, 31, 32, 42, 49, 61, 62, 67, 74, 76, 83, 84, 85, 90, 115, 149, 156, 162, 167, 169, 172], "ar": [0, 18, 20, 21, 22, 23, 25, 26, 28, 29, 31, 32, 38, 39, 42, 47, 48, 49, 61, 64, 73, 74, 76, 108, 115, 123, 145, 146, 147, 148, 162, 167, 168, 171, 172], "gan": [0, 2, 6, 60, 72, 111, 112, 114, 123, 168, 171, 172], "instance_seg": [0, 2, 6, 72, 75], "labelfre": [0, 6, 168, 171, 172], "segmen": 0, "when": [0, 25, 26, 28, 29, 30, 31, 61, 76, 138, 168, 171], "experi": [0, 5, 172], "overrid": [0, 7, 168, 171], "list": [0, 7, 11, 13, 18, 20, 22, 25, 28, 29, 31, 32, 38, 40, 42, 43, 45, 50, 51, 53, 57, 62, 67, 74, 76, 79, 81, 90, 116, 145, 146, 147, 148, 149, 158, 167], "string": [0, 13, 18, 20, 21, 31, 61, 138], "format": [0, 38, 57, 172], "same": [0, 21, 57, 62, 90, 115, 123, 172], "wai": [0, 168, 171], "cli": 0, "load_default_experi": [0, 2, 7], "batch_siz": [0, 25, 26, 28, 29, 30, 31, 84, 160, 168, 171, 172], "16": [0, 6, 147, 148, 149], "onc": [0, 172], "i": [0, 4, 18, 20, 21, 22, 23, 25, 26, 28, 29, 30, 31, 32, 33, 42, 49, 51, 52, 53, 61, 62, 76, 79, 81, 83, 84, 85, 90, 115, 123, 126, 138, 144, 149, 158, 160, 163, 167, 168, 171, 172], "attribut": 0, "overridden": [0, 172], "override_config": [0, 2, 7], "which": [0, 7, 20, 21, 25, 26, 28, 29, 31, 32, 40, 61, 77, 107, 163, 167, 168, 171, 172], "take": [0, 18, 51, 53], "valu": [0, 6, 20, 21, 22, 23, 25, 28, 32, 38, 47, 48, 49, 79, 81, 83, 84, 85, 128, 138, 143, 158, 160, 167, 168, 171, 172], "made": [0, 28, 61, 172], "async": 0, "pass": [0, 25, 26, 28, 29, 31, 32, 33, 38, 39, 51, 61, 64, 67, 73, 74, 77, 83, 84, 85, 90, 107, 168, 171, 172], "run_async": [0, 7], "true": [0, 7, 23, 28, 30, 32, 38, 39, 42, 45, 47, 51, 59, 76, 87, 90, 110, 111, 113, 114, 116, 128, 138, 145, 146, 147, 148, 149, 150, 168, 171], "await": 0, "The": [0, 20, 23, 28, 29, 32, 49, 61, 62, 83, 84, 85, 123, 138, 163, 168, 171, 172], "To": [0, 168, 171], "examin": [0, 46], "print_config": [0, 2, 7], "download": 0, "exampl": [0, 32, 51, 61, 77, 115, 138, 168, 171], "download_example_data": [0, 2, 7], "thi": [0, 4, 14, 23, 28, 29, 30, 32, 49, 51, 61, 62, 77, 160, 168, 171, 172], "api": [1, 170], "cyto_dl_model": [1, 2], "modul": [1, 2, 3, 8, 17, 19, 24, 27, 36, 37, 44, 58, 60, 63, 66, 68, 72, 75, 78, 82, 86, 89, 95, 99, 105, 106, 109, 118, 133, 141, 142, 151, 152, 155, 164, 170], "autothresholdmethod": [1, 2, 6], "experimenttyp": [1, 2, 6], "hardwaretyp": [1, 2, 5, 6], "patchsiz": [1, 2, 5, 6], "callback": [1, 162, 167, 168, 170, 171], "callback_util": [1, 8], "csv_saver": [1, 8], "csvsaver": [1, 8, 10], "jepasav": [1, 8, 10], "image_sav": [1, 8], "imagesav": [1, 8, 11], "latent_walk": [1, 8], "layer_freez": [1, 8], "layerfreez": [1, 8, 13], "model_util": [1, 8], "getembed": [1, 8, 14], "get_all_embed": [1, 8, 14], "save_predictions_classifi": [1, 8, 14], "outlier_detect": [1, 8], "outlierdetect": [1, 8, 15], "datafram": [1, 24, 33, 168, 170, 171], "transform": [1, 17, 25, 26, 28, 29, 30, 31, 32, 33, 36, 38, 42, 43, 76, 79, 115, 140, 146, 147, 148, 168, 171, 172], "reader": [1, 17], "filter_column": [1, 17, 18, 19, 20], "read_csv": [1, 17, 18], "read_datafram": [1, 17, 18], "read_h5ad": [1, 17, 18], "read_parquet": [1, 17, 18], "datamodul": [1, 14, 168, 170, 171], "arrai": [1, 24, 32, 42, 81, 115, 126, 155, 168, 171], "make_array_dataload": [1, 24, 25], "data_dict": [1, 24], "make_data_dict_dataload": [1, 24, 26], "folder": [1, 24, 162, 168, 171], "make_folder_dataload": [1, 24, 31], "multidim_imag": [1, 24], "multidimimagedataset": [1, 24, 32], "make_multidim_image_dataload": [1, 24, 32], "smartcach": [1, 24], "smartcachedatamodul": [1, 24, 33], "torchvis": [1, 24], "cifar10": [1, 24, 34], "cifar100": [1, 24, 34], "clevrclassif": [1, 24, 34], "caltech101": [1, 24, 34], "caltech256": [1, 24, 34], "celeba": [1, 24, 34], "cityscap": [1, 24, 34], "cococapt": [1, 24, 34], "cocodetect": [1, 24, 34], "country211": [1, 24, 34], "dtd": [1, 24, 34], "datasetfold": [1, 24, 34], "emnist": [1, 24, 34], "eurosat": [1, 24, 34], "fer2013": [1, 24, 34], "fgvcaircraft": [1, 24, 34], "fakedata": [1, 24, 34], "fashionmnist": [1, 24, 34], "flickr30k": [1, 24, 34], "flickr8k": [1, 24, 34], "flowers102": [1, 24, 34], "food101": [1, 24, 34], "gtsrb": [1, 24, 34], "hmdb51": [1, 24, 34], "inaturalist": [1, 24, 34], "imagefold": [1, 24, 34], "imagenet": [1, 24, 34], "kmnist": [1, 24, 34], "kinet": [1, 24, 34], "kitti": [1, 24, 34], "lfwpair": [1, 24, 34], "lfwpeopl": [1, 24, 34], "lsun": [1, 24, 34], "lsunclass": [1, 24, 34], "mnist": [1, 24, 34], "movingmnist": [1, 24, 34], "omniglot": [1, 24, 34], "oxfordiiitpet": [1, 24, 34], "pcam": [1, 24, 34], "phototour": [1, 24, 34], "places365": [1, 24, 34], "qmnist": [1, 24, 34], "renderedsst2": [1, 24, 34], "sbdataset": [1, 24, 34], "sbu": [1, 24, 34], "semeion": [1, 24, 34], "stl10": [1, 24, 34], "sun397": [1, 24, 34], "svhn": [1, 24, 34], "stanfordcar": [1, 24, 34], "ucf101": [1, 24, 34], "usp": [1, 24, 34], "vocdetect": [1, 24, 34], "vocsegment": [1, 24, 34], "visiondataset": [1, 24, 34], "widerfac": [1, 24, 34], "imag": [1, 11, 26, 29, 30, 32, 33, 64, 67, 73, 74, 76, 79, 81, 83, 84, 108, 110, 111, 112, 113, 114, 115, 116, 123, 126, 140, 145, 146, 147, 148, 168, 170, 171, 172], "io": [1, 18, 36, 73, 74, 115, 146, 151, 168, 171], "logger": [1, 161, 162, 167, 168, 170, 171], "mlflow": [1, 58, 60, 86, 168, 171], "mlflowlogg": [1, 58, 59], "classif": [1, 60], "contrast": [1, 48, 60], "handler": [1, 60], "im2im": [1, 60], "jepa": [1, 49, 60, 105, 141], "util": [1, 24, 25, 26, 27, 28, 29, 31, 60, 72, 105, 141, 168, 170, 171, 172], "vae": [1, 60], "base_model": [1, 60], "basemodel": [1, 60, 61, 62, 64, 67, 73, 74, 83, 84, 85, 90], "basemodelmeta": [1, 60, 61], "basic_model": [1, 60], "basicmodel": [1, 60, 62], "nn": [1, 62, 64, 67, 73, 74, 79, 83, 84, 85, 90, 170], "discrimin": [1, 73, 77, 105, 111, 119, 123, 168, 171], "head": [1, 64, 67, 73, 74, 105, 146, 147, 148, 149, 156, 168, 171, 172], "loss": [1, 61, 62, 76, 90, 102, 105, 110, 111, 112, 113, 114, 115, 116, 168, 171, 172], "point_cloud": [1, 105, 170], "vit": [1, 105], "hr_skip": [1, 105], "hrskip": [1, 105, 117], "mlp": [1, 105, 116, 141, 142, 143], "res_unit": [1, 105], "residualunit": [1, 105, 138], "spatial_transform": [1, 105], "convpoolrelu": [1, 105, 139], "stn": [1, 105, 139], "track_sequence_predictor": [1, 105], "trackclassifi": [1, 105, 140], "positionalencoding1d": [1, 105, 140], "spharm": [1, 155], "create_dataload": [1, 155, 156], "extract_array_predict": [1, 155, 156], "checkpoint": [1, 4, 15, 87, 149, 155], "load_checkpoint": [1, 155, 157], "config": [1, 4, 7, 155, 156, 162, 167, 168, 171], "is_config": [1, 155, 158], "kv_to_dict": [1, 155, 158], "remove_aux_kei": [1, 155, 158], "download_test_data": [1, 155], "delete_test_data": [1, 155, 159], "parse_s3_path": [1, 155, 159], "setup_path": [1, 155, 159], "dummy_dataset": [1, 155], "dummydatamodul": [1, 155, 160], "dummydataset": [1, 155, 160], "pylogg": [1, 155], "get_pylogg": [1, 155, 161], "rich_util": [1, 155], "enforce_tag": [1, 155, 162], "print_config_tre": [1, 155, 162], "rotat": [1, 155, 164, 172], "rotationmodul": [1, 155, 163], "get_rotation_matrix_so2": [1, 155, 163], "get_rotation_matrix_so3": [1, 155, 163], "template_util": [1, 155], "close_logg": [1, 155, 167], "extra": [1, 145, 147, 148, 155, 162, 167, 169], "get_metric_valu": [1, 155, 167], "instantiate_callback": [1, 155, 167], "instantiate_logg": [1, 155, 167], "log_hyperparamet": [1, 155, 167], "save_fil": [1, 155, 167], "task_wrapp": [1, 155, 167], "compil": [1, 64, 73, 74, 168, 170, 171], "eval": [1, 170], "main": [1, 35, 120, 154, 170], "train": [1, 2, 3, 4, 5, 6, 7, 11, 23, 25, 26, 28, 29, 30, 31, 32, 49, 51, 59, 61, 62, 64, 73, 74, 76, 83, 84, 85, 87, 114, 149, 168, 170, 171], "cyto_dl_base_model": [2, 3], "cytodlbasemodel": [2, 3, 4, 5], "segmentation_plugin_model": [2, 3], "segmentationpluginmodel": [2, 3, 5], "isodata": [2, 6], "li": [2, 6], "local": [2, 6], "mean": [2, 6, 46, 76, 101, 119, 121, 122], "minimum": [2, 6, 33], "multiotsu": [2, 6], "niblack": [2, 6], "otsu": [2, 6], "sauvola": [2, 6], "triangl": [2, 6], "yen": [2, 6], "label_fre": [2, 6], "segment": [2, 6, 52, 76, 77, 149, 168, 171, 172], "segmentation_arrai": [2, 6], "segmentation_plugin": [2, 5, 6], "cpu": [2, 6, 61, 81], "gpu": [2, 6, 161, 167, 172], "larg": [2, 6, 33, 172], "medium": [2, 6], "small": [2, 6, 76], "load_config_from_fil": [2, 7], "predict": [2, 3, 4, 5, 7, 10, 25, 28, 29, 32, 61, 62, 64, 67, 83, 84, 85, 110, 111, 112, 113, 114, 115, 123, 126, 147, 156, 168, 171], "save_config": [2, 3, 4, 7], "from_default_config": [3, 4], "from_existing_config": [3, 4], "get_config": [3, 4, 60, 86, 87], "get_experiment_nam": [3, 4], "set_experiment_nam": [3, 4], "get_hardware_typ": [3, 5], "get_input_channel": [3, 5], "get_manifest_column_nam": [3, 5], "get_patch_s": [3, 5], "get_raw_image_channel": [3, 5], "get_split_column": [3, 5], "remove_split_column": [3, 5], "set_hardware_typ": [3, 5], "set_input_channel": [3, 5], "set_manifest_column_nam": [3, 5], "set_patch_s": [3, 5], "set_raw_image_channel": [3, 5], "set_split_column": [3, 5], "dictconfig": [4, 5, 26, 35, 154, 158, 162, 167], "sourc": [4, 5, 6, 7, 10, 11, 13, 14, 15, 18, 20, 21, 22, 23, 25, 26, 28, 29, 30, 31, 32, 33, 35, 38, 39, 40, 41, 42, 43, 45, 46, 47, 48, 49, 50, 51, 52, 53, 55, 57, 59, 61, 62, 64, 65, 67, 73, 74, 76, 77, 79, 80, 81, 83, 84, 85, 87, 88, 90, 100, 101, 102, 103, 107, 108, 110, 111, 112, 113, 114, 115, 116, 117, 119, 120, 121, 122, 123, 124, 126, 128, 129, 130, 131, 132, 138, 139, 140, 143, 144, 145, 146, 147, 148, 149, 150, 154, 156, 157, 158, 159, 160, 161, 162, 163, 167], "base": [4, 5, 6, 7, 10, 11, 13, 14, 15, 21, 22, 26, 28, 29, 30, 31, 32, 33, 34, 38, 39, 40, 41, 42, 43, 45, 46, 47, 48, 49, 50, 51, 52, 53, 55, 57, 59, 61, 62, 64, 65, 67, 73, 74, 76, 77, 79, 80, 81, 83, 84, 85, 90, 100, 101, 102, 103, 107, 108, 110, 111, 112, 113, 114, 115, 116, 117, 119, 120, 121, 122, 123, 124, 126, 128, 129, 130, 131, 132, 138, 139, 140, 143, 144, 145, 146, 147, 148, 149, 160, 163, 168, 171, 172], "abc": [4, 110], "A": [4, 5, 18, 20, 22, 26, 28, 29, 31, 62, 123, 138, 163], "us": [4, 5, 14, 18, 20, 21, 22, 23, 25, 26, 28, 29, 31, 32, 33, 38, 39, 47, 48, 49, 50, 51, 52, 61, 62, 64, 67, 73, 74, 76, 77, 83, 84, 85, 90, 110, 112, 114, 115, 123, 138, 145, 146, 147, 148, 160, 162, 168, 169, 171, 172], "run": [4, 32, 74, 110, 113, 114, 116, 168, 171, 172], "cyto": [4, 169], "dl": [4, 169], "Not": [4, 114], "intend": 4, "direct": [4, 139], "client": 4, "pleas": 4, "see": [4, 18, 20, 25, 26, 28, 29, 31, 115, 138, 168, 171], "classmethod": [4, 101], "factori": 4, "method": [4, 61, 80, 115], "instead": [4, 38, 84], "spatial_dim": [4, 32, 33, 112, 115, 138, 145, 146, 147, 148, 149, 163, 172], "int": [4, 5, 7, 11, 22, 23, 28, 29, 30, 32, 33, 46, 48, 49, 51, 52, 53, 57, 67, 76, 77, 79, 83, 84, 85, 100, 102, 107, 108, 112, 115, 116, 126, 132, 138, 145, 146, 147, 148, 149, 150, 160, 163], "return": [4, 18, 23, 32, 51, 61, 76, 107, 123, 140, 144, 158, 163], "paramet": [4, 7, 11, 13, 18, 20, 21, 22, 23, 25, 26, 28, 29, 30, 31, 32, 33, 38, 40, 42, 45, 46, 47, 48, 49, 50, 51, 52, 53, 55, 57, 62, 64, 67, 73, 74, 76, 77, 79, 81, 83, 84, 85, 90, 101, 107, 108, 110, 111, 112, 113, 114, 115, 116, 123, 126, 128, 140, 145, 146, 147, 148, 149, 158, 163, 167, 168, 171, 172], "dimens": [4, 32, 33, 49, 76, 115, 116, 138, 140, 144, 145, 146, 147, 148, 149, 150], "e": [4, 18, 32, 90, 147, 168, 169, 171, 172], "g": [4, 32, 147, 168, 171, 172], "2": [4, 29, 32, 51, 76, 107, 108, 117, 138, 139, 140, 148, 149, 163], "config_filepath": 4, "an": [4, 6, 18, 25, 26, 29, 32, 49, 51, 53, 61, 79, 101, 126, 138, 168, 171, 172], "exist": [4, 23], "basi": 4, "must": [4, 13, 18, 20, 23, 25, 31, 32, 45, 51, 61, 168, 171], "gener": [4, 18, 22, 33, 49, 51, 57, 62, 76, 77, 79, 83, 84, 85, 108, 110, 111, 112, 150, 160, 168, 171], "subclass": [4, 30, 90], "want": [4, 168, 171, 172], "str": [4, 5, 7, 11, 13, 14, 15, 18, 20, 21, 22, 23, 25, 26, 28, 29, 30, 31, 32, 33, 38, 39, 40, 41, 42, 43, 45, 46, 47, 48, 50, 51, 52, 53, 55, 57, 59, 62, 64, 67, 73, 74, 76, 79, 80, 81, 83, 84, 85, 90, 114, 123, 126, 138, 145, 147, 149, 162, 167], "manifest_path": 4, "output_dir": [4, 7, 14], "none": [4, 5, 7, 11, 13, 14, 18, 20, 21, 22, 23, 26, 28, 29, 30, 31, 32, 33, 35, 41, 43, 51, 52, 59, 61, 62, 64, 74, 76, 79, 80, 90, 101, 108, 110, 111, 112, 113, 114, 115, 116, 117, 124, 131, 132, 138, 143, 145, 146, 147, 148, 149, 154, 156, 160, 162, 167], "max_epoch": [4, 83, 84, 85], "handl": [5, 32, 42, 61, 123], "tupl": [5, 32, 39, 41, 43, 49, 61, 79, 138, 146, 148], "hardware_typ": 5, "input_channel": [5, 145, 147, 148], "target1": 5, "target2": 5, "merge_mask": 5, "exclude_mask": 5, "base_imag": 5, "patch_siz": [5, 140, 145, 147, 172], "channel": [5, 32, 33, 38, 41, 43, 46, 52, 57, 79, 108, 115, 126, 138, 145, 147, 148, 149, 168, 171, 172], "split_column": [5, 28, 29, 30, 168, 171], "enum": 6, "enumer": [6, 33, 38], "threshold_isodata": 6, "threshold_li": 6, "threshold_loc": 6, "threshold_mean": 6, "threshold_minimum": 6, "threshold_multiotsu": 6, "threshold_niblack": 6, "threshold_otsu": 6, "threshold_sauvola": 6, "threshold_triangl": 6, "threshold_yen": 6, "patch": [6, 46, 49, 51, 145, 146, 147, 148, 149, 150, 172], "size": [6, 30, 49, 51, 76, 123, 138, 145, 146, 147, 148, 149, 172], "respect": [6, 18, 90, 172], "shape": [6, 42, 160], "128": [6, 140], "64": [6, 108, 139], "32": [6, 139, 148, 149], "object": [7, 18, 32, 46, 76, 77, 79, 80, 81, 108, 123, 149, 163], "dict": [7, 20, 21, 25, 26, 28, 29, 32, 39, 41, 43, 51, 52, 59, 62, 73, 74, 76, 83, 84, 85, 90, 160, 167], "load": [7, 18, 26, 28, 29, 31, 32, 33, 38, 39, 40, 62, 168, 171, 172], "config_path": 7, "experiment_typ": 7, "directori": [7, 11, 26, 28, 29, 31, 64, 67, 73, 74, 84, 172], "float": [7, 23, 33, 46, 47, 48, 49, 76, 77, 79, 80, 83, 84, 85, 90, 123, 126, 128, 138, 148, 154, 163, 167], "bool": [7, 11, 20, 23, 28, 29, 30, 32, 38, 39, 40, 41, 42, 43, 45, 47, 48, 51, 52, 53, 55, 57, 76, 79, 81, 84, 112, 115, 123, 126, 128, 138, 145, 146, 147, 148, 149, 150, 162], "fals": [7, 11, 20, 23, 28, 29, 30, 32, 38, 39, 40, 45, 47, 48, 51, 53, 55, 57, 64, 73, 74, 76, 79, 81, 84, 90, 100, 101, 103, 108, 112, 114, 115, 117, 119, 138, 143, 145, 146, 147, 148, 162], "save": [7, 8, 11, 32, 36, 44, 57, 62, 64, 67, 73, 74, 84, 110, 113, 114, 116, 167, 168, 171, 172], "current": [7, 29, 32, 42], "on_predict_epoch_end": [8, 10], "pred_to_csv": [8, 10], "save_feat": [8, 10], "on_predict_batch_end": [8, 11, 15], "on_test_batch_end": [8, 11, 15], "on_train_batch_end": [8, 11], "on_validation_batch_end": [8, 11], "setup": [8, 13, 24, 33, 167], "on_test_epoch_end": [8, 14], "calculate_mahalanobis_hook": [8, 15], "flatten_activ": [8, 15], "on_load_checkpoint": [8, 15], "on_predict_epoch_start": [8, 15], "on_save_checkpoint": [8, 15], "on_test_epoch_start": [8, 15], "on_train_epoch_start": [8, 15], "update_covariance_hook": [8, 15], "save_dir": [10, 11, 15, 59, 64, 67, 73, 74, 84, 172], "meta_kei": [10, 45, 67], "trainer": [10, 11, 13, 14, 15, 162, 168, 171], "pl_modul": [10, 11, 13, 14, 15], "pred": [10, 14, 61, 107, 120], "stage": [10, 11, 13, 33, 61, 62, 64, 67, 73, 74, 83, 84, 85, 90, 110, 111, 113, 114, 116], "save_every_n_epoch": 11, "1": [11, 23, 29, 30, 33, 49, 51, 53, 64, 73, 74, 76, 79, 90, 101, 110, 111, 112, 113, 115, 116, 123, 126, 129, 131, 132, 138, 145, 147, 148, 159, 172], "val": [11, 23, 28, 29, 61, 62, 88], "save_input": 11, "after": [11, 45, 76, 79, 115, 143, 167], "postprocess": [11, 72, 75, 110, 111, 112, 113, 114, 115, 116, 168, 171, 172], "ead": 11, "union": [11, 13, 18, 20, 21, 25, 26, 28, 29, 31, 32, 33, 40, 42, 50, 53, 62, 74, 76, 81], "frequenc": [11, 64, 67, 73, 74], "whether": [11, 23, 28, 29, 30, 32, 38, 39, 40, 42, 46, 47, 48, 51, 52, 53, 57, 64, 73, 74, 76, 79, 81, 84, 114, 115, 128, 138, 148, 160, 162, 163, 172], "input": [11, 20, 22, 23, 25, 26, 28, 29, 31, 40, 62, 64, 73, 74, 76, 83, 85, 108, 110, 111, 112, 113, 114, 115, 116, 119, 121, 122, 123, 124, 131, 138, 145, 147, 148, 158, 162, 168, 171, 172], "output": [11, 15, 32, 38, 62, 79, 90, 110, 115, 116, 123, 138, 148, 149, 156, 162, 167, 168, 171], "batch": [11, 15, 21, 26, 28, 29, 30, 31, 32, 45, 55, 61, 62, 64, 65, 67, 73, 74, 83, 84, 85, 90, 101, 110, 111, 113, 114, 116, 160, 163, 172], "batch_idx": [11, 15, 61, 62, 64, 65, 67, 73, 74, 83, 84, 85, 90], "dataloader_idx": [11, 15], "0": [11, 15, 33, 41, 46, 47, 48, 49, 51, 76, 77, 83, 84, 85, 90, 101, 112, 115, 117, 123, 126, 128, 138, 143, 145, 147, 148, 149, 163], "step": [11, 32, 59, 61, 77, 172], "contain": [13, 18, 20, 21, 25, 26, 28, 31, 32, 33, 42, 172], "exclud": [13, 18, 20, 31, 76], "search": 13, "within": [13, 48], "layer": [13, 108, 115, 116, 123, 138, 144, 146, 147, 148, 149], "freez": 13, "x_label": [14, 62, 90], "id_label": [14, 61, 90], "arg": [14, 25, 26, 28, 29, 31, 34, 38, 40, 61, 62, 64, 65, 67, 73, 74, 83, 84, 85, 90, 138, 160, 162], "id_field": 14, "lightningmodul": [14, 61, 167], "train_dataload": [14, 24, 27, 28, 33, 155, 160], "val_dataload": [14, 24, 27, 28, 33, 155, 160], "test_dataload": [14, 24, 27, 28, 33, 155, 160], "todo": [14, 23, 28, 29], "better": [14, 145, 146, 147, 148, 172], "mayb": 14, "vol": 14, "predictor": [14, 83, 84, 85, 147], "code": [14, 168, 171], "drop": [14, 30, 143], "unnecessari": 14, "index": [14, 53, 150], "n_epoch": 15, "layer_nam": 15, "act": [15, 138], "set": [15, 28, 48, 167], "forward": [15, 60, 61, 62, 63, 64, 66, 67, 72, 74, 82, 85, 89, 90, 99, 100, 101, 102, 103, 105, 106, 107, 108, 109, 110, 111, 112, 115, 117, 118, 119, 120, 121, 122, 124, 129, 130, 131, 132, 138, 139, 140, 141, 142, 143, 144, 145, 146, 147, 148, 149], "hook": 15, "callabl": [15, 25, 26, 31, 32, 51, 62, 79, 112, 115, 167], "filter": [17, 18, 19, 31, 108, 149], "filter_row": [17, 19, 20], "group_col": [17, 19], "groupcol": [17, 19, 21], "misc": [17, 19], "append_class_weight": [17, 19, 22], "append_labels_to_integ": [17, 19, 22], "append_one_hot": [17, 19, 22], "make_random_df": [17, 19, 22], "split": [17, 19, 28, 29, 30, 33, 61, 168, 171, 172], "sample_n_each": [17, 19, 23], "split_datafram": [17, 19, 23], "columns_to_filt": 18, "sequenc": [18, 20, 22, 25, 26, 28, 29, 30, 31, 39, 40, 41, 43, 45, 51, 55, 62, 76, 79, 90, 138, 140, 145, 150, 162, 168, 171], "regex": [18, 20, 31], "startswith": [18, 20, 31], "endswith": [18, 20, 31], "column": [18, 20, 21, 22, 23, 28, 29, 30, 31, 32, 33, 168, 171, 172], "combin": [18, 30, 148, 168, 171, 172], "differ": [18, 20, 23, 38, 123, 126, 147, 149, 168, 171], "queri": [18, 143], "pattern": 18, "If": [18, 20, 22, 25, 28, 29, 30, 31, 32, 38, 39, 45, 62, 115, 145, 146, 147, 148, 168, 171, 172], "suppli": [18, 20], "preced": 18, "remain": [18, 20, 21], "argument": [18, 22, 25, 26, 28, 29, 31, 33, 39, 64, 67, 73, 74, 83, 84, 85, 90, 107, 138, 168, 171], "ignor": [18, 20, 28, 29, 42, 167], "otherwis": [18, 62, 172], "logic": [18, 61], "AND": 18, "appli": [18, 25, 26, 31, 33, 51, 52, 53, 57, 79, 115, 144, 167, 168, 171], "all": [18, 23, 28, 29, 32, 33, 51, 74, 79, 167, 168, 171, 172], "condit": 18, "regular": [18, 20, 31], "express": [18, 20, 31], "match": [18, 20, 31, 61], "substr": [18, 20, 31], "start": [18, 20, 31, 32, 33, 51, 57, 167, 168, 171], "end": [18, 20, 31], "include_column": 18, "read": [18, 32, 33, 37, 38, 39, 41, 43], "store": [18, 26, 28, 29, 31, 62, 168, 171], "csv": [18, 29, 32, 33, 38, 67, 168, 171, 172], "includ": [18, 20, 23, 145, 147, 148], "onli": [18, 20, 25, 26, 28, 29, 30, 31, 32, 42, 52, 62, 67, 76, 110, 138, 146, 148, 167, 172], "given": [18, 20, 22, 23, 28, 29, 46, 61, 102, 128, 150, 168, 171], "upath": [18, 26, 28, 29, 31], "desir": 18, "result": [18, 20, 61, 62, 172], "type": [18, 21, 29, 30, 32, 38, 61, 79, 112, 115, 123, 138, 158, 163, 168, 171, 172], "pd": [18, 20, 22, 23], "required_column": 18, "parquet": [18, 28, 29], "assert": 18, "expect": [18, 26, 168, 171, 172], "requir": [18, 168, 169, 171, 172], "either": [18, 20, 21, 23, 28, 61, 150, 163], "support": [18, 29, 42, 77, 123], "aren": 18, "t": [18, 33, 38, 102, 169], "found": [18, 168, 171], "valueerror": 18, "thrown": 18, "those": [18, 168, 171], "get": [18, 33, 168, 171], "append": [18, 22, 23], "without": [18, 150], "duplic": 18, "back": 18, "anndata": 18, "h5ad": 18, "r": [18, 169], "": [18, 20, 28, 29, 30, 73, 74, 76, 77, 90, 111, 112, 163, 168, 171], "doc": [18, 25, 26, 28, 29, 31, 73, 74, 115, 168, 171], "detail": [18, 20], "http": [18, 25, 26, 28, 29, 31, 39, 43, 73, 74, 115, 120, 121, 146, 168, 169, 171], "readthedoc": 18, "en": [18, 73, 74, 115, 168, 171], "latest": 18, "html": [18, 25, 26, 28, 29, 31, 73, 74, 115, 168, 171], "select": [20, 23, 30, 46, 168, 171], "dataset": [20, 25, 28, 29, 30, 32, 33, 160, 168, 171], "more": [20, 25, 26, 28, 29, 31, 167, 168, 171], "oper": 20, "It": [20, 28, 101, 172], "panda": [20, 23], "case": [20, 21, 23, 28, 29, 61, 168, 171, 172], "explicit": [20, 130], "keep": [20, 76, 168, 171], "row": [20, 22, 23], "where": [20, 21, 29, 32, 61, 79, 146, 149, 158, 160, 168, 171], "group": [21, 29, 163], "dtype": [21, 38, 79, 126], "kei": [21, 26, 28, 29, 30, 31, 32, 38, 40, 42, 45, 46, 47, 48, 50, 51, 52, 53, 55, 57, 61, 62, 64, 67, 73, 74, 76, 81, 83, 84, 85, 90, 143, 158, 160, 168, 171], "becom": [21, 160, 168, 171], "kwarg": [21, 33, 34, 38, 40, 61, 62, 64, 65, 67, 73, 74, 83, 84, 85, 90, 100, 101, 102, 103, 107, 120, 160], "col": 21, "filter_col": 21, "singl": [21, 25, 28, 32, 168, 171], "np": [21, 25, 79], "numpi": [21, 25, 32, 38, 42, 79, 81, 156], "recogn": 21, "turn": 21, "add": [22, 45, 61], "weight": [22, 62, 76, 83, 84, 85, 90, 111, 112, 126, 131, 144, 149], "modifi": [22, 23, 29, 107], "its": [22, 46, 162], "integ": [22, 32, 51], "encod": [22, 49, 83, 84, 85, 89, 90, 140, 143, 145, 146, 147, 148], "convert": [22, 32, 42, 76, 81, 145, 156], "hot": 22, "call": [22, 28, 29, 61, 77, 167], "function": [22, 32, 51, 62, 76, 90, 110, 111, 112, 113, 114, 115, 116, 126, 128, 167, 168, 171, 172], "one_hot_encod": 22, "b": [22, 130], "c": [22, 38], "d": 22, "n_row": 22, "100": [22, 51, 76, 83, 84, 85, 111, 112], "random": [22, 23, 28, 29, 51], "test": [22, 23, 28, 29, 32, 61, 62, 168, 171], "wrangl": 22, "pipelin": 22, "creat": [22, 25, 26, 28, 29, 31, 32, 42, 64, 76, 123, 149, 156, 168, 169, 171, 172], "number": [22, 23, 28, 29, 33, 38, 49, 51, 52, 57, 77, 83, 84, 85, 107, 108, 115, 126, 138, 145, 146, 147, 148, 149, 150, 167, 172], "forc": 23, "seed": [23, 28, 29], "42": [23, 28, 29], "have": [23, 45, 51, 61, 76, 138], "equal": [23, 126], "per": [23, 28, 29, 51, 76, 126, 156], "In": [23, 61, 163, 168, 171, 172], "ha": [23, 77, 123], "less": 23, "than": [23, 138], "correspond": [23, 29, 31, 51, 160], "sampl": [23, 25, 26, 28, 29, 30, 31, 46, 49, 51, 99, 100, 101, 103, 160], "replac": [23, 33], "uniqu": [23, 168, 171], "toggl": [23, 105, 139], "upsampl": [23, 115, 149, 168, 171], "smaller": 23, "train_frac": 23, "val_frac": 23, "return_split": 23, "perform": [23, 115], "three": [23, 28], "identifi": 23, "each": [23, 25, 26, 28, 29, 30, 31, 33, 49, 57, 61, 64, 76, 90, 138, 140, 144, 145, 146, 147, 148, 149, 160, 168, 171, 172], "belong": 23, "extend": 23, "enabl": 23, "balanc": 23, "stratifi": 23, "fraction": 23, "valid": [23, 33, 61, 64, 73, 74, 168, 171], "By": [23, 172], "half": 23, "between": [23, 29, 30, 32, 64, 73, 74, 101, 110, 115, 126, 147], "separ": [23, 32, 38, 76], "reproduc": 23, "dataframe_datamodul": [24, 27], "dataframedatamodul": [24, 27, 28, 29], "grouped_dataframe_datamodul": [24, 27], "groupeddataframedatamodul": [24, 27, 29], "alternatingbatchsampl": [24, 27, 29, 30], "removenankeysd": [24, 27, 30], "get_canonical_split_nam": [24, 27, 30], "get_dataset": [24, 27, 28, 30], "make_multiple_dataframe_split": [24, 27, 30], "make_single_dataframe_split": [24, 27, 30], "parse_transform": [24, 27, 30], "create_metatensor": [24, 32], "get_per_file_arg": [24, 32, 33], "is_batch": [24, 32], "make_dataload": [24, 27, 28, 29, 33], "predict_dataload": [24, 27, 28, 33, 155, 160], "prepare_data": [24, 33], "ndarrai": [25, 39, 41, 43, 76, 81], "source_kei": [25, 84], "dataloader_kwarg": [25, 26, 28, 29, 31, 32], "dataload": [25, 26, 28, 29, 31, 32, 33, 62, 156, 168, 171], "element": [25, 61, 138, 160], "keyword": [25, 26, 28, 29, 31, 39], "torch": [25, 26, 28, 29, 30, 31, 42, 62, 64, 73, 74, 79, 90, 101, 108, 111, 112, 115, 123, 143, 144, 163, 172], "instanti": [25, 26, 28, 29, 31, 90, 167, 168, 171], "asid": [25, 26, 28, 29, 31], "shuffl": [25, 26, 28, 29, 30, 31], "among": [25, 26, 28, 29, 31], "num_work": [25, 26, 28, 29, 31, 33, 168, 171], "etc": [25, 26, 28, 29, 31, 168, 171, 172], "pytorch": [25, 26, 28, 29, 30, 31, 62, 90, 168, 171], "info": [25, 26, 28, 29, 31], "org": [25, 26, 28, 29, 31, 121], "stabl": [25, 26, 28, 29, 31, 73, 74, 115, 168, 171], "cache_dir": [26, 28, 29, 30, 31, 168, 171], "usual": 26, "cach": [26, 28, 29, 31, 33, 168, 171], "acceler": [26, 28, 29, 31], "get_dataload": [27, 28, 155, 160], "split_map": [28, 29, 30], "just_infer": [28, 29, 30], "subsampl": [28, 29], "refresh_subsampl": [28, 29], "smartcache_arg": [28, 29, 30], "lightningdatamodul": [28, 33, 160], "lightn": [28, 62, 167, 168, 171, 172], "fold": 28, "addition": [28, 167], "go": 28, "flag": 28, "so": [28, 29, 77, 163, 168, 171, 172], "whole": [28, 29], "predict_datamodul": 28, "simpli": [28, 61, 168, 171], "construct": [28, 168, 171], "regardless": 28, "specif": [28, 29, 61, 73, 74, 90, 115, 168, 171, 172], "everyth": [28, 29, 48], "document": [28, 29, 168, 171], "just": [28, 29, 84], "infer": [28, 29, 61, 73, 74, 90, 100, 101, 103, 163], "epoch": [28, 29, 83, 84, 85], "refresh": [28, 29], "time": [28, 29, 30, 77, 167], "smartcachedataset": [28, 29], "target_column": [29, 30], "grouping_column": [29, 30], "should": [29, 32, 33, 40, 51, 61, 115, 126, 138, 149, 172], "some": [29, 32, 138, 168, 171, 172], "criterion": 29, "leverag": 29, "two": [29, 50, 101, 168, 171], "multitask": 29, "ground": [29, 30, 76, 115, 123, 168, 171, 172], "truth": [29, 30, 76, 115, 123, 168, 171, 172], "task": [29, 30, 64, 67, 73, 74, 77, 110, 111, 113, 114, 115, 116, 145, 167, 168, 171, 172], "characterist": 29, "altern": [29, 30, 172], "dure": [29, 32, 64, 67, 73, 74, 114, 149, 167, 168, 171], "factor": [29, 51, 149, 172], "homogen": 29, "across": [29, 53, 76, 168, 171], "subset": 30, "drop_last": 30, "sampler": 30, "subsetrandomsampl": 30, "batchsampl": 30, "mutual": 30, "exclus": 30, "monai": [30, 33, 51, 53, 73, 74, 115, 138, 168, 171, 172], "wrap": [30, 62, 167, 168, 171], "repres": [30, 138], "last": [30, 52, 87, 123, 138, 144], "incomplet": 30, "randomli": [30, 49, 52, 138], "follow": [30, 172], "order": [30, 138, 145, 148, 149, 162], "remov": [30, 49, 76], "nan": 30, "ad": [30, 45, 49, 172], "allow_missing_kei": [30, 38, 47, 48, 51, 52, 53, 55, 57, 76, 81], "allow": [30, 38, 52, 55, 57, 172], "multi": [30, 32, 161, 167, 168, 171, 172], "target": [30, 49, 61, 67, 77, 84, 119, 121, 122, 123, 124, 128, 130, 131, 147, 168, 171, 172], "split_path": 30, "dataframe_path": 30, "orig_fnam": 31, "origin": [31, 45], "filenam": [31, 39, 41, 43], "extens": 31, "out": [31, 64, 73, 74, 138, 148, 149, 168, 171], "rule": [31, 172], "csv_path": [32, 33], "img_path_column": [32, 33], "channel_column": [32, 33], "out_kei": [32, 38], "3": [32, 33, 51, 76, 77, 107, 108, 112, 115, 117, 138, 145, 146, 147, 148, 149, 163, 169, 172], "scene_column": 32, "scene": [32, 33, 38, 39], "time_start_column": 32, "time_stop_column": 32, "stop": 32, "time_step_column": 32, "dict_meta": 32, "dask_load": [32, 38, 39], "dimension": [32, 100, 101, 102], "timelaps": 32, "metadata": [32, 45, 160], "timepoint": [32, 33, 38, 39], "extract": [32, 38, 39, 83, 84, 85, 148], "spatial": [32, 33, 76, 107, 115, 126, 138, 145, 147, 148, 149, 172], "yx": [32, 49, 145, 172], "zyx": [32, 49, 145, 148, 149, 172], "specifi": [32, 33, 38, 39, 48, 53, 81, 90, 168, 171, 172], "multipl": [32, 51, 107, 108, 138, 147, 168, 171], "thei": [32, 45, 77], "comma": [32, 38], "scene1": 32, "scene2": 32, "ani": [32, 59, 79, 126, 168, 171, 172], "start_column": 32, "stop_column": 32, "step_column": 32, "For": [32, 51, 115, 138, 168, 171, 172], "everi": [32, 62, 158], "other": [32, 168, 171, 172], "version": [32, 107, 172], "accept": [32, 126], "normal": [32, 108, 138, 143, 144, 168, 171, 172], "could": 32, "here": [32, 33, 61, 121, 168, 171, 172], "dask": [32, 38, 39], "full": [32, 38, 39], "memori": [32, 33, 38, 39, 172], "befor": [32, 38, 39, 51, 57, 128, 167], "img": [32, 39, 41, 43, 57, 76, 81, 140, 145, 147, 148, 149], "meta": 32, "df": [32, 33], "x": [32, 61, 62, 64, 74, 85, 90, 107, 108, 110, 111, 112, 115, 117, 120, 129, 132, 138, 139, 143, 144, 172], "listconfig": [32, 52, 57, 158], "multidimimag": 32, "cannot": 32, "dim": [32, 76, 107, 108, 116, 132], "multiscen": 32, "alreadi": [32, 168, 171], "compos": [33, 162], "img_data": 33, "n_val": 33, "20": 33, "pct_val": 33, "raw": [33, 38, 51, 168, 171, 172], "ch": [33, 46, 79], "num_neighbor": 33, "4": [33, 140, 143, 146, 148, 149], "cache_r": 33, "5": [33, 49, 51, 76, 129], "replace_r": 33, "czi": 33, "don": 33, "fit": 33, "bioio": 33, "n_imag": 33, "percentag": 33, "neighbor": [33, 76], "worker": 33, "most": [33, 61], "schedul": 33, "parallel": 33, "submodul": [36, 151, 170], "aicsimage_load": [36, 37], "aicsimageloaderd": [36, 37, 38], "monai_bio_read": [36, 37, 38], "monaibioread": [36, 37, 39], "numpy_read": [36, 37], "readnumpyfil": [36, 37, 40], "ome_zarr_read": [36, 37], "omezarrread": [36, 37, 41], "polygon_load": [36, 37], "polygonloaderd": [36, 37, 42], "skimage_read": [36, 37], "skimageread": [36, 37, 43], "add_meta": [36, 44], "addmeta": [36, 44, 45], "bright_sampl": [36, 44], "brightsampl": [36, 44, 46], "clip": [36, 44, 48], "clipd": [36, 44, 47], "contrastadjust": [36, 44], "contrastadjustd": [36, 44, 48], "generate_jepa_mask": [36, 44], "jepamaskgener": [36, 44, 49], "merg": [36, 44], "multiscale_cropp": [36, 44], "randommultiscalecropd": [36, 44, 51], "pad": [36, 44, 76, 138], "padzd": [36, 44, 52], "project": [36, 44, 168, 171], "maxprojectd": [36, 44, 53], "rotation_mask_transform": [36, 44], "so2_random_rot": [36, 44], "track_transform": [36, 44], "cropres": [36, 44, 57], "generatetracklabel": [36, 44, 57], "perchannel": [36, 44, 57], "splittrackd": [36, 44, 57], "split_arg": [37, 38], "get_data": [37, 39, 41, 43], "verify_suffix": [37, 39, 41, 43], "path_kei": 38, "scene_kei": 38, "kwargs_kei": 38, "dimension_order_out": [38, 172], "float16": 38, "channel_kei": 38, "timepoint_kei": 38, "fix": [38, 145, 146, 147, 148, 150], "initi": [38, 77, 115, 123, 161], "bioimag": [38, 39], "get_image_dask_data": [38, 39], "miss": [38, 42, 47, 48, 52, 53, 55, 57, 76, 81], "adapt": [39, 43, 120], "github": [39, 43, 120, 146, 169], "com": [39, 43, 120, 169], "mmv": [39, 43], "lab": [39, 43], "mmv_im2im": [39, 43], "reader_kwarg": 39, "imageread": [39, 41, 43], "get_image_data": 39, "pathlik": [39, 41, 43], "maptransform": 40, "thereof": 40, "interpret": 40, "point": [40, 51, 76, 138, 172], "cloud": [40, 120], "remot": 40, "fsspec": 40, "locat": [40, 168, 171], "level": [41, 168, 171, 172], "image_nam": 41, "shape_reference_kei": 42, "propagate_3d": 42, "missing_key_mod": 42, "rais": [42, 53, 76, 81], "polygon": 42, "vertic": 42, "tensor": [42, 81, 84, 85, 101, 123, 124, 131, 138, 149, 163, 172], "mask": [42, 49, 50, 85, 114, 145, 146, 148, 149], "propag": 42, "2d": [42, 49, 83, 115, 138, 145, 163, 168, 171, 172], "3d": [42, 49, 76, 83, 115, 138, 145, 163, 168, 171, 172], "how": [42, 115, 160, 168, 171], "keyerror": 42, "do": [42, 115, 123, 168, 171], "noth": 42, "new": [42, 45, 168, 171, 172], "blank": 42, "remove_excess_pixel": [44, 49], "generate_slic": [44, 51], "image_kei": [45, 50, 52, 76], "delet": 45, "updat": [45, 61], "metatensor": [45, 160], "been": 45, "threshold": [46, 79, 128], "base_prob": 46, "decid": [46, 51], "intens": [46, 47, 48, 79, 172], "background": [46, 48], "rate": [46, 138], "probabilii": 46, "check": [46, 52], "against": 46, "low": [47, 48, 49], "01": 47, "high": [47, 48, 49], "99": 47, "percentil": [47, 48], "absolut": [47, 48], "lower": [47, 48], "bound": [47, 48], "upper": [47, 48], "per_channel": 47, "resiz": [47, 48, 51, 57, 172], "fail": [47, 48], "adjust": 48, "rang": [48, 49, 79, 168, 171, 172], "outsid": 48, "everywher": 48, "mask_siz": 49, "12": [49, 147, 148, 149], "block_aspect_ratio": 49, "num_patch": [49, 145, 146, 147, 148, 149, 150], "6": [49, 76, 149], "24": 49, "mask_ratio": [49, 145, 148, 149], "9": [49, 169], "randomizabletransform": [49, 51, 57, 76], "block": [49, 105, 115, 138, 141], "contigu": 49, "work": [49, 101, 145, 146, 147, 148, 169, 172], "until": [49, 172], "met": 49, "exceed": [49, 172], "exterior": 49, "exactli": 49, "determin": [49, 115, 138, 162, 168, 171, 172], "aspect": 49, "ratio": [49, 115, 138, 148, 149], "proport": 49, "pixel": [49, 57, 76, 126, 145, 147, 148, 172], "along": 49, "boundari": 49, "reach": 49, "mask_kei": [50, 114], "base_image_kei": 50, "output_nam": 50, "serv": 50, "patch_shap": [51, 139, 172], "scales_dict": 51, "patch_per_imag": 51, "selection_fn": 51, "max_attempt": 51, "style": [51, 53, 148, 172], "slice": [51, 52, 76, 172], "scale": [51, 107, 111, 112, 123, 126], "superresolut": 51, "x_kei": [51, 64, 73, 74, 83, 85], "network": [51, 62, 64, 73, 74, 90, 115, 138, 172], "Its": 51, "resolut": [51, 112, 115], "len": 51, "map": [51, 110, 138, 172], "seg": [51, 76, 105, 141, 168, 171], "would": [51, 61], "kept": 51, "discard": 51, "max": [51, 53], "attempt": 51, "try": [51, 172], "quit": 51, "image_dict": 51, "segmentation_kei": 52, "pad_amount": 52, "pad_kei": 52, "segmentation_ch": 52, "top": [52, 149, 172], "bottom": 52, "crop": 52, "repeat": 52, "first": [52, 101, 108, 163, 168, 171, 172], "present": [52, 76], "offset": 52, "presenc": 52, "projection_dim": 53, "nczyx": 53, "comput": [53, 101, 128, 163], "error": [53, 76, 81], "save_path": 55, "disk": 55, "debug": 55, "max_shift": 57, "8": [57, 101, 140, 143, 146, 148, 149, 168, 171], "maximum": [57, 83, 84, 85], "shift": 57, "img_kei": 57, "formation_kei": 57, "breakdown_kei": 57, "breakdown": 57, "track_start_kei": 57, "track_start": 57, "label_kei": [57, 76], "label": [57, 64, 67, 76, 79, 123], "track": 57, "after_save_checkpoint": [58, 59], "log_hyperparam": [58, 59], "log_metr": [58, 59], "experiment_nam": [59, 168, 171], "lightning_log": 59, "run_nam": [59, 168, 171], "tracking_uri": [59, 87], "tag": [59, 147, 162, 167], "mlrun": 59, "prefix": 59, "artifact_loc": 59, "run_id": [59, 87], "fault_toler": 59, "ckpt_callback": 59, "param": [59, 110], "namespac": 59, "mode": [59, 87, 100, 101, 102, 103, 117, 121, 130, 167], "metric": [59, 61, 110, 113, 114, 116, 167], "timepoint_classif": [60, 63], "timepointclassif": [60, 63, 65], "base_handl": [60, 68], "image_handl": [60, 68], "load_image_patch": [60, 68], "multi_task": [60, 72], "multitaskim2im": [60, 72, 73, 74], "ijepa": [60, 82, 85], "iwm": [60, 82, 85], "jepa_bas": [60, 82], "jepabas": [60, 82, 83, 84, 85], "load_model_from_checkpoint": [60, 86, 87], "find_indic": [60, 86, 88], "o2_spharm_va": [60, 89], "prior": [60, 89, 90], "base_va": [60, 89], "baseva": [60, 89, 90], "image_canon_va": [60, 89], "image_encod": [60, 89], "image_va": [60, 89], "latent_loss_va": [60, 89], "point_cloud_va": [60, 89], "tabular_va": [60, 89], "compute_metr": [60, 61], "configure_optim": [60, 61, 72, 73, 74, 82, 85], "model_step": [60, 61, 62, 63, 64, 66, 67, 72, 73, 74, 82, 83, 84, 85, 89, 90], "on_train_start": [60, 61], "parse_batch": [60, 61], "predict_step": [60, 61, 63, 64, 65, 66, 67, 72, 73, 74, 82, 84, 85], "test_step": [60, 61, 72, 74, 82, 84], "training_step": [60, 61, 72, 74], "validation_step": [60, 61, 72, 74], "log": [61, 62, 101, 167, 168, 171, 172], "assumpt": 61, "_step": 61, "whose": 61, "mai": 61, "self": [61, 76, 83, 84], "structur": [61, 162, 172], "part": [61, 90, 167, 168, 171], "arbitrari": 61, "denot": 61, "you": [61, 168, 171, 172], "implement": [61, 168, 171, 172], "process": [61, 167, 172], "variabl": [61, 168, 171], "hparam": 61, "id": 61, "detach": [61, 75, 78, 81, 110, 111, 112, 113, 114, 115], "consist": [61, 149, 168, 171], "your": [61, 168, 171, 172], "might": 61, "wish": 61, "post": 61, "minim": 62, "wrapper": [62, 126, 128], "around": [62, 145, 147, 148], "assum": [62, 163, 168, 171, 172], "gt": [62, 120], "optim": [62, 111, 112, 168, 171, 172], "retriev": [62, 167], "adam": 62, "save_predict": 62, "serotini": [62, 168, 171], "fields_to_log": 62, "field": [62, 162, 172], "pretrained_weight": 62, "pretrain": [62, 149], "via": 62, "load_state_dict": 62, "save_imag": [63, 64], "should_save_imag": [63, 64], "share": [64, 73, 74, 172], "save_images_every_n_epoch": [64, 73, 74], "base_kwarg": [64, 67, 73, 74, 83, 84, 85, 90], "logit": [64, 115, 168, 171], "text": 64, "overlaid": 64, "plot_class": [66, 67], "plot_neighbor": [66, 67], "backbon": [67, 73, 74, 110, 113, 114, 115, 116, 149, 168, 171, 172], "task_head": [67, 73, 74, 156, 172], "anchor_kei": 67, "anchor": 67, "positive_kei": 67, "posit": [67, 140, 145, 146, 147, 148, 150], "target_kei": [67, 84, 90], "visual": [67, 168, 171], "viz_freq": 67, "x1": 67, "x2": 67, "embedding1": 67, "embedding2": 67, "instancesegclust": [72, 75, 76], "instancesegloss": [72, 75, 76], "instancesegpreprocessd": [72, 75, 76], "instancesegrandflipd": [72, 75, 76], "pad_slic": [72, 75, 76], "noise_anneal": [72, 75, 108], "noiseanneal": [72, 75, 77], "get_n_postprocess_imag": [72, 74], "get_per_head": [72, 74], "run_forward": [72, 74, 110, 111, 113, 114, 116], "basic": [73, 90, 168, 171, 172], "inference_arg": [73, 74], "slide": [73, 74], "window": [73, 74], "inference_head": 74, "run_head": [74, 109, 110, 111, 112, 113, 114, 116], "n_postprocess": [74, 110, 111, 113, 114, 116], "act_thresh_label": [75, 78], "actthreshlabel": [75, 78, 79], "auto_thresh": [75, 78], "autothreshold": [75, 78, 80], "dict_to_im": [75, 78], "dicttoim": [75, 78, 81], "cluster_object": [75, 76], "kd_cluster": [75, 76], "remove_small_skeleton": [75, 76], "embed_from_skel": [75, 76], "keep_largest_cc": [75, 76], "label_2d": [75, 76], "shrink": [75, 76], "skeleton_tal": [75, 76], "smooth_embed": [75, 76], "topology_preserving_thin": [75, 76], "update_nois": [75, 77], "anisotropi": 76, "skel_threshold": 76, "semantic_threshold": 76, "min_siz": 76, "1000": 76, "distance_threshold": 76, "progress": 76, "cluster": 76, "instanceseg": 76, "find": 76, "skeleton": 76, "assign": 76, "semant": 76, "embed": [76, 84, 145, 146, 147, 148, 149, 150], "nearest": [76, 117], "distanc": [76, 172], "skel": 76, "closest": 76, "below": [76, 128, 168, 171], "touch": 76, "edg": [76, 172], "compon": [76, 162], "kernel_s": [76, 138, 139], "thin": 76, "keep_largest": 76, "instanc": [76, 77, 138, 149, 172], "kernel": [76, 138], "gaussian": [76, 77, 89, 99], "smooth": 76, "flow": 76, "amount": [76, 115, 138], "psuedo": 76, "largest": 76, "connect": [76, 115, 138, 149], "iseg": 76, "vector": [76, 163, 172], "z": [76, 90, 100, 101, 102, 163, 168, 171], "im": [76, 108], "topologi": 76, "preserv": 76, "binari": 76, "max_label": 76, "increas": [76, 77, 172], "thick": 76, "convolv": 76, "non": 76, "bw": 76, "bridg": 76, "gap": 76, "eros": 76, "spatial_axi": 76, "prob": [76, 121], "flip": [76, 172], "augment": [76, 168, 171, 172], "sign": 76, "gradient": 76, "chang": [76, 168, 171, 172], "axi": [76, 163], "NOT": 76, "probabl": 76, "constraint": 76, "annealing_step": 77, "5000": 77, "init_vari": 77, "anneal": [77, 108], "varianc": [77, 101], "nois": [77, 108], "real": [77, 107, 123], "fake": [77, 123], "harder": 77, "distribut": [77, 102], "overlap": 77, "nice": 77, "theoret": 77, "implic": 77, "qualiti": 77, "also": [77, 101, 138, 168, 171, 172], "curriculum": 77, "learn": [77, 83, 84, 172], "techniqu": 77, "iter": 77, "unblur": 77, "over": 77, "linearli": 77, "activ": [79, 107, 115, 138, 143, 169], "ident": [79, 112, 115, 138, 168, 171], "_type": 79, "_dtype_lik": 79, "_supportsdtyp": 79, "supportsindex": 79, "_dtypedict": 79, "uint8": 79, "rescale_dtyp": 79, "purpos": 79, "rescal": [79, 168, 171], "dtypelik": 79, "multichannel": 81, "concaten": 81, "cuda": 81, "extract_embed": [82, 84], "get_predict_mask": [82, 84], "get_context_embed": [82, 85], "get_mask": [82, 85, 142, 145], "get_target_embed": [82, 85], "remove_first_dim": [82, 85], "update_teach": [82, 85], "supervis": [83, 84], "featur": [83, 84, 85, 110, 113, 114, 116, 123, 138, 143, 144, 146, 147, 148, 149, 168, 171], "access": [83, 84, 85], "momentum": [83, 84, 85], "exponenti": [83, 84, 85], "move": [83, 84, 85], "averag": [83, 84, 85, 126], "998": [83, 84, 85], "world": 84, "translat": 84, "latent": [84, 90], "space": [84, 90], "target_domain_kei": 84, "domain": [84, 147], "predict_target": 84, "devic": [84, 145], "tmp_dir": 87, "strict": 87, "ckpt": [87, 172], "lst": 88, "o2_spharm_encod": [89, 95], "abstract_prior": [89, 99], "diagonalgaussianprior": [89, 99, 101], "isotropicgaussianprior": [89, 99, 101], "compute_tc_penalti": [89, 99, 101], "identity_prior": [89, 99], "identityprior": [89, 99, 102], "joint_prior": [89, 99], "jointprior": [89, 99, 103], "calculate_elbo": [89, 90], "calculate_rcl": [89, 90], "calculate_rcl_dict": [89, 90], "decod": [89, 90, 143, 146, 148, 149], "sample_z": [89, 90], "beta": [90, 130], "kld": 90, "term": [90, 138], "reconstruction_loss": [90, 111, 112], "reconstruct": [90, 111, 112, 155, 164], "_loss": [90, 119, 121, 122, 124, 130, 131], "abstractprior": 90, "decoder_latent_part": 90, "decoer": 90, "what": [90, 162], "depend": [90, 169], "xhat": 90, "input_kei": 90, "return_param": 90, "z_parts_param": 90, "kl_diverg": [99, 100, 101, 103], "param_s": [99, 100, 101, 102, 103], "kl": [100, 101, 102, 103], "properti": [100, 101, 102, 103], "logvar": 101, "learn_mean": 101, "learn_logvar": 101, "clamp_logvar": 101, "tc_penalty_weight": 101, "mu1": 101, "mu2": 101, "logvar1": 101, "logvar2": 101, "reduct": [101, 103, 119, 121, 122, 130, 131], "sum": [101, 103, 130, 144, 149], "kullback": 101, "leibler": 101, "diverg": 101, "diagon": 101, "necessarili": 101, "isotrop": 101, "wise": 101, "second": 101, "doesn": [102, 169], "contribut": 102, "effect": 102, "dirac": 102, "delta": 102, "z_param": 103, "multi_scale_discrimin": [105, 106], "multiscalediscrimin": [105, 106, 107], "n_layer_discrimin": [105, 106], "nlayerdiscrimin": [105, 106, 107, 108], "base_head": [105, 109], "basehead": [105, 109, 110, 111, 113, 114, 115, 116], "gan_head": [105, 109], "ganhead": [105, 109, 111, 112], "gan_head_superr": [105, 109], "ganhead_res": [105, 109, 112], "mae_head": [105, 109], "maehead": [105, 109, 113], "mask_head": [105, 109], "maskhead": [105, 109, 114], "res_blocks_head": [105, 109], "resblockshead": [105, 109, 112, 115], "vic_reg": [105, 109, 118], "vicreghead": [105, 109, 116], "adversarial_loss": [105, 118], "adversarialloss": [105, 118, 119], "chamfer_loss": [105, 118], "chamferloss": [105, 118, 120], "continuous_bernoulli": [105, 118], "cblogloss": [105, 118, 121], "cosine_loss": [105, 118], "cosineloss": [105, 118, 122], "gan_loss": [105, 111, 112, 118], "ganloss": [105, 111, 112, 118, 123], "pix2pixhd": [105, 107, 111, 112, 118, 123], "gaussian_nll_loss": [105, 118], "gaussiannllloss": [105, 118, 124], "geomloss": [105, 118], "loss_wrapp": [105, 118], "cmap_loss": [105, 118, 126], "losswrapp": [105, 118, 126], "spharm_loss": [105, 118], "threshold_loss": [105, 118], "thresholdloss": [105, 118, 128], "vicregloss": [105, 118, 129], "off_diagon": [105, 118, 129], "weibul": [105, 118], "weibulllogloss": [105, 118, 130], "weibull_log_prob": [105, 118, 130], "weighted_mse_loss": [105, 118], "weightedmseloss": [105, 118, 131], "dgcnn": [105, 133], "folding_net": [105, 133], "graph_funct": [105, 133], "vnn": [105, 133], "cross_ma": [105, 141], "crossmae_decod": [105, 141, 146], "iwmpredictor": [105, 141, 147], "jepaencod": [105, 141, 147], "jepapredictor": [105, 141, 147], "mae": [105, 141], "mae_decod": [105, 141, 148], "mae_encod": [105, 141, 148], "mae_vit": [105, 141, 148], "encodedskip": [105, 141, 149], "seg_vit": [105, 141, 149], "superresdecod": [105, 141, 149], "get_positional_embed": [105, 141, 150], "take_index": [105, 141, 150], "n_scale": 107, "input_nc": 108, "ndf": 108, "n_layer": 108, "norm_lay": [108, 143, 144], "instancenorm": 108, "instancenorm3d": 108, "defin": [108, 123], "patchgan": 108, "n_channel": 108, "input_im": 108, "model_output": 108, "conv_fn": 108, "later": 108, "requires_featur": 108, "standard": 108, "generate_io_map": [109, 110], "update_param": [109, 110], "input_filenam": 110, "backbone_featur": [110, 111, 113, 114, 116], "y_hat": [110, 111, 113, 114, 116], "calcul": [110, 113, 114, 116, 126], "bcewithlogitsloss": [111, 112, 123], "feature_matching_loss": [111, 112], "l1loss": [111, 112], "mseloss": [111, 112], "reconstruction_loss_weight": [111, 112], "in_channel": [112, 115, 117, 138, 172], "out_channel": [112, 115, 117, 138, 149, 172], "final_act": [112, 115], "lr": [112, 115], "n_conv": [112, 115], "dropout": [112, 115, 117, 138], "upsample_method": [112, 115], "pixelshuffl": [112, 115], "upsample_ratio": [112, 115], "first_lay": [112, 115], "dens": [112, 115], "inherit": 112, "__init__": 112, "calculate_metr": 114, "convolut": [115, 138, 145, 147, 148, 149, 168, 171], "final": [115, 148], "done": 115, "hr": 115, "length": [115, 126, 140], "convprojectionlay": 115, "2048": 116, "8192": 116, "hidden_lay": [116, 132], "projectornetwork": 116, "hidden": 116, "scale_factor": 117, "align_corn": 117, "antialia": 117, "batch_pairwise_dist": [118, 120], "batch_pairwise_dist2": [118, 120], "get_target_tensor": [118, 123], "get_feature_matching_loss": [118, 123], "get_gan_loss": [118, 123], "argmax": 119, "squeez": 119, "return_pr": 119, "sentinal4d": 120, "cellshap": 120, "blob": 120, "cellshape_cloud": 120, "vendor": 120, "chamfer_dist": 120, "py": [120, 172], "licens": 120, "license_antao": 120, "y": [120, 129, 143], "continu": 121, "bernoulli": 121, "propos": 121, "arxiv": 121, "ab": 121, "1907": 121, "06845": 121, "gan_mod": 123, "vanilla": 123, "target_real_label": 123, "target_fake_label": 123, "abstract": [123, 168, 171], "awai": 123, "need": [123, 168, 171, 172], "lsgan": 123, "wgangp": 123, "note": [123, 172], "sigmoid": [123, 168, 171], "target_is_r": 123, "fill": 123, "loss_weight": 123, "fm": 123, "10": [123, 124, 130, 138], "feature_typ": 123, "mean_dim": 124, "ep": [124, 130, 138, 163], "1e": [124, 130, 138, 163], "costmap": 126, "differenti": [126, 168, 171], "emphas": 126, "throughout": 126, "loss_fn": [126, 128], "channel_weight": 126, "_supportsarrai": 126, "_nestedsequ": 126, "complex": [126, 168, 171], "byte": 126, "loss_scal": 126, "arraylik": 126, "abov": 128, "num_featur": 129, "sim_coeff": 129, "std_coeff": 129, "cov_coeff": 129, "alpha": 130, "256": [132, 139], "scale_output": 132, "stride": 138, "subunit": 138, "adn_ord": 138, "nda": 138, "prelu": 138, "norm": [138, 144], "dropout_dim": 138, "dilat": 138, "bia": 138, "last_conv_onli": 138, "residu": 138, "conv": 138, "init": 138, "normalized_shap": 138, "print": [138, 162, 167], "sequenti": 138, "unit0": 138, "conv3d": 138, "adn": 138, "num_paramet": 138, "n": [138, 168, 169, 171], "layernorm": [138, 143, 144], "05": 138, "elementwise_affin": 138, "unit1": 138, "zero": [138, 167], "entir": 138, "larger": 138, "control": [138, 167], "implicit": 138, "both": [138, 172], "side": 138, "in_filt": 139, "out_filt": 139, "n_input_ch": 139, "512": 139, "n_conv_filt": 139, "pos_embedding_length": 140, "120": 140, "emb_dim": [140, 143, 145, 146, 147, 148, 149, 150], "num_lay": [140, 144, 146, 147, 148, 149], "num_head": [140, 143, 146, 147, 148], "num_class": 140, "token": [140, 146, 150], "classifi": 140, "them": [140, 168, 171], "d_model": 140, "matrix": [140, 163], "cross_attent": [141, 142], "crossattent": [141, 142, 143], "crossattentionblock": [141, 142, 143], "crossselfblock": [141, 142, 143], "intermediate_weigh": [141, 142], "intermediateweigh": [141, 142, 144], "patchifi": [141, 142, 147], "random_index": [141, 142, 145], "init_weight": [141, 146, 147, 148], "predict_target_featur": [141, 147], "encoder_dim": 143, "decoder_dim": [143, 148], "qkv_bia": 143, "qk_scale": 143, "attn_drop": 143, "proj_drop": 143, "mlp_ratio": 143, "drop_path": 143, "act_lay": 143, "gelu": 143, "in_featur": 143, "hidden_featur": 143, "out_featur": 143, "embed_dim": 144, "n_output": 144, "intermedi": [144, 148, 149], "n_patch": 145, "context_pixel": [145, 147, 148], "learnable_pos_embed": [145, 146, 147, 148], "pix": 145, "learnabl": [145, 146, 147, 148, 150], "sin": [145, 146, 147, 148, 163], "co": [145, 146, 147, 148, 163], "empir": [145, 146, 147, 148], "brightfield": [145, 146, 147, 148], "n_visible_patch": 145, "base_patch_s": [146, 148, 149], "enc_dim": [146, 148], "768": [146, 148, 149], "192": [146, 147, 148, 149], "inspir": 146, "crossma": [146, 148], "attend": 146, "visibl": 146, "forward_index": [146, 148], "backward_index": [146, 148], "input_dim": 147, "special": 147, "condition": 147, "fluoresc": 147, "context_emb": 147, "target_mask": 147, "target_domain": 147, "context": 147, "n_intermediate_weight": 148, "weight_intermedi": 148, "linear": 148, "like": [148, 149, 172], "75": [148, 149], "encoder_lay": [148, 149], "encoder_head": [148, 149], "decoder_lay": [148, 149], "decoder_head": 148, "use_crossma": 148, "features_onli": 148, "n_decoder_filt": 149, "upsample_factor": 149, "6134": 149, "5005": 149, "encoder_ckpt": 149, "freeze_encod": 149, "encoder_kwarg": 149, "simpl": [149, 168, 171], "aic": 149, "20x": 149, "100x": 149, "unet": 149, "fed": [149, 168, 171], "skip": 149, "use_cls_token": 150, "cl": 150, "read_pcloud": [151, 152], "generate_img": [155, 160], "compute_rotation_matrix": [155, 163], "data_cfg": 156, "load_param": 157, "kv": 158, "item": 158, "nest": 158, "form": [158, 163], "interpol": [158, 168, 171], "limit": [159, 168, 171, 172], "fn": 159, "num_sampl": 160, "dummy_metadata": 160, "10000": 160, "mani": [160, 168, 171, 172], "k": 160, "friendli": 161, "save_to_fil": 162, "prompt": 162, "user": [162, 168, 171], "print_ord": 162, "resolv": 162, "content": [162, 167], "rich": [162, 167], "librari": 162, "tree": 162, "hydra": [162, 168, 171, 172], "refer": [162, 168, 171], "export": 162, "pad_valu": 163, "06": 163, "pose": 163, "about": 163, "equivari": [163, 169], "theta": 163, "indic": 163, "pair": 163, "precis": 163, "sure": 167, "close": 167, "properli": 167, "prevent": 167, "failur": 167, "multirun": 167, "warn": 167, "metric_dict": 167, "metric_nam": 167, "safe": 167, "callbacks_cfg": 167, "logger_cfg": 167, "object_dict": 167, "rank": 167, "task_func": 167, "decor": 167, "resist": 167, "finish": 167, "except": 167, "occur": 167, "total": 167, "execut": 167, "dir": 167, "aim": [168, 171], "common": [168, 171], "manner": [168, 171], "easi": [168, 171], "beginn": [168, 171], "flexibl": [168, 171, 172], "custom": [168, 171], "templat": [168, 171, 172], "conveni": [168, 171], "syntax": [168, 171], "repo": [168, 171], "best": [168, 171], "familiar": [168, 171, 172], "mechan": [168, 171], "topic": [168, 171], "cover": [168, 171], "becaus": [168, 171], "help": [168, 171], "These": [168, 171], "reli": [168, 171], "heavili": [168, 171], "omegaconf": [168, 171], "expos": [168, 171], "_aux": [168, 171], "section": [168, 171], "persist": [168, 171], "speed": [168, 171], "up": [168, 171, 172], "divid": [168, 171], "manifest": [168, 171], "preprocess": [168, 171], "offer": [168, 171], "wide": [168, 171], "often": [168, 171, 172], "primari": [168, 171], "due": [168, 171], "score": [168, 171], "v": [168, 171], "binar": [168, 171, 172], "own": [168, 171], "thing": [168, 171], "pin_memori": [168, 171], "workhors": [168, 171], "source_col": [168, 171, 172], "target_col": [168, 171], "organ": [168, 171], "arrang": [168, 171], "produc": [168, 171], "through": [168, 171], "further": [168, 171], "lr_schedul": [168, 171], "alwai": [168, 171], "well": [168, 171], "_target_": [168, 171], "_partial_": [168, 171], "sigmoid_rescal": [168, 171], "tell": [168, 171], "bit": [168, 171], "encourag": [168, 171], "few": [168, 171, 172], "git": 169, "clone": 169, "allencellmodel": 169, "cd": 169, "conda": 169, "environ": 169, "myenv": 169, "pip": 169, "txt": 169, "relat": 169, "equiv": 169, "upgrad": 169, "subpackag": 170, "outlin": 172, "modif": 172, "doe": 172, "suit": 172, "aicsimageio": 172, "recommend": 172, "suitabl": 172, "robust": 172, "again": 172, "excel": 172, "ensur": 172, "while": 172, "naiv": 172, "incorrect": 172, "neural": 172, "architectur": 172, "signal": 172, "cut": 172, "crucial": 172, "bottleneck": 172, "suggest": 172, "look": 172, "consider": 172, "primarili": 172, "As": 172, "thumb": 172, "enough": 172, "entiti": 172, "establish": 172, "too": 172, "downsampl": 172, "decreas": 172, "usag": 172, "ckpt_path": 172, "boolean": 172, "dictat": 172, "omit": 172}, "objects": {"": [[1, 0, 0, "-", "cyto_dl"]], "cyto_dl": [[2, 0, 0, "-", "api"], [8, 0, 0, "-", "callbacks"], [17, 0, 0, "-", "dataframe"], [24, 0, 0, "-", "datamodules"], [35, 0, 0, "-", "eval"], [36, 0, 0, "-", "image"], [58, 0, 0, "-", "loggers"], [60, 0, 0, "-", "models"], [105, 0, 0, "-", "nn"], [151, 0, 0, "-", "point_cloud"], [154, 0, 0, "-", "train"], [155, 0, 0, "-", "utils"]], "cyto_dl.api": [[3, 0, 0, "-", "cyto_dl_model"], [6, 0, 0, "-", "data"], [7, 0, 0, "-", "model"]], "cyto_dl.api.cyto_dl_model": [[4, 0, 0, "-", "cyto_dl_base_model"], [5, 0, 0, "-", "segmentation_plugin_model"]], "cyto_dl.api.cyto_dl_model.cyto_dl_base_model": [[4, 1, 1, "", "CytoDLBaseModel"]], "cyto_dl.api.cyto_dl_model.cyto_dl_base_model.CytoDLBaseModel": [[4, 2, 1, "", "from_default_config"], [4, 2, 1, "", "from_existing_config"], [4, 2, 1, "", "get_config"], [4, 2, 1, "", "get_experiment_name"], [4, 2, 1, "", "predict"], [4, 2, 1, "", "save_config"], [4, 2, 1, "", "set_experiment_name"], [4, 2, 1, "", "train"]], "cyto_dl.api.cyto_dl_model.segmentation_plugin_model": [[5, 1, 1, "", "SegmentationPluginModel"]], "cyto_dl.api.cyto_dl_model.segmentation_plugin_model.SegmentationPluginModel": [[5, 2, 1, "", "get_hardware_type"], [5, 2, 1, "", "get_input_channel"], [5, 2, 1, "", "get_manifest_column_names"], [5, 2, 1, "", "get_patch_size"], [5, 2, 1, "", "get_raw_image_channels"], [5, 2, 1, "", "get_split_column"], [5, 2, 1, "", "remove_split_column"], [5, 2, 1, "", "set_hardware_type"], [5, 2, 1, "", "set_input_channel"], [5, 2, 1, "", "set_manifest_column_names"], [5, 2, 1, "", "set_patch_size"], [5, 2, 1, "", "set_raw_image_channels"], [5, 2, 1, "", "set_split_column"]], "cyto_dl.api.data": [[6, 1, 1, "", "AutoThresholdMethod"], [6, 1, 1, "", "ExperimentType"], [6, 1, 1, "", "HardwareType"], [6, 1, 1, "", "PatchSize"]], "cyto_dl.api.data.AutoThresholdMethod": [[6, 3, 1, "", "ISODATA"], [6, 3, 1, "", "LI"], [6, 3, 1, "", "LOCAL"], [6, 3, 1, "", "MEAN"], [6, 3, 1, "", "MINIMUM"], [6, 3, 1, "", "MULTIOTSU"], [6, 3, 1, "", "NIBLACK"], [6, 3, 1, "", "OTSU"], [6, 3, 1, "", "SAUVOLA"], [6, 3, 1, "", "TRIANGLE"], [6, 3, 1, "", "YEN"]], "cyto_dl.api.data.ExperimentType": [[6, 3, 1, "", "GAN"], [6, 3, 1, "", "INSTANCE_SEG"], [6, 3, 1, "", "LABEL_FREE"], [6, 3, 1, "", "SEGMENTATION"], [6, 3, 1, "", "SEGMENTATION_ARRAY"], [6, 3, 1, "", "SEGMENTATION_PLUGIN"]], "cyto_dl.api.data.HardwareType": [[6, 3, 1, "", "CPU"], [6, 3, 1, "", "GPU"]], "cyto_dl.api.data.PatchSize": [[6, 3, 1, "", "LARGE"], [6, 3, 1, "", "MEDIUM"], [6, 3, 1, "", "SMALL"]], "cyto_dl.api.model": [[7, 1, 1, "", "CytoDLModel"]], "cyto_dl.api.model.CytoDLModel": [[7, 2, 1, "", "download_example_data"], [7, 2, 1, "", "load_config_from_dict"], [7, 2, 1, "", "load_config_from_file"], [7, 2, 1, "", "load_default_experiment"], [7, 2, 1, "", "override_config"], [7, 2, 1, "", "predict"], [7, 2, 1, "", "print_config"], [7, 2, 1, "", "save_config"], [7, 2, 1, "", "train"]], "cyto_dl.callbacks": [[10, 0, 0, "-", "csv_saver"], [11, 0, 0, "-", "image_saver"], [13, 0, 0, "-", "layer_freeze"], [14, 0, 0, "-", "model_utils"], [15, 0, 0, "-", "outlier_detection"]], "cyto_dl.callbacks.csv_saver": [[10, 1, 1, "", "CSVSaver"], [10, 1, 1, "", "JEPASaver"]], "cyto_dl.callbacks.csv_saver.CSVSaver": [[10, 2, 1, "", "on_predict_epoch_end"], [10, 2, 1, "", "pred_to_csv"], [10, 2, 1, "", "save_feats"]], "cyto_dl.callbacks.csv_saver.JEPASaver": [[10, 2, 1, "", "pred_to_csv"]], "cyto_dl.callbacks.image_saver": [[11, 1, 1, "", "ImageSaver"]], "cyto_dl.callbacks.image_saver.ImageSaver": [[11, 2, 1, "", "on_predict_batch_end"], [11, 2, 1, "", "on_test_batch_end"], [11, 2, 1, "", "on_train_batch_end"], [11, 2, 1, "", "on_validation_batch_end"], [11, 2, 1, "", "save"]], "cyto_dl.callbacks.layer_freeze": [[13, 1, 1, "", "LayerFreeze"]], "cyto_dl.callbacks.layer_freeze.LayerFreeze": [[13, 2, 1, "", "setup"]], "cyto_dl.callbacks.model_utils": [[14, 1, 1, "", "GetEmbeddings"], [14, 4, 1, "", "get_all_embeddings"], [14, 4, 1, "", "save_predictions_classifier"]], "cyto_dl.callbacks.model_utils.GetEmbeddings": [[14, 2, 1, "", "on_test_epoch_end"]], "cyto_dl.callbacks.outlier_detection": [[15, 1, 1, "", "OutlierDetection"]], "cyto_dl.callbacks.outlier_detection.OutlierDetection": [[15, 2, 1, "", "calculate_mahalanobis_hook"], [15, 2, 1, "", "flatten_activations"], [15, 2, 1, "", "on_load_checkpoint"], [15, 2, 1, "", "on_predict_batch_end"], [15, 2, 1, "", "on_predict_epoch_start"], [15, 2, 1, "", "on_save_checkpoint"], [15, 2, 1, "", "on_test_batch_end"], [15, 2, 1, "", "on_test_epoch_start"], [15, 2, 1, "", "on_train_epoch_start"], [15, 2, 1, "", "update_covariance_hook"]], "cyto_dl.dataframe": [[18, 0, 0, "-", "readers"], [19, 0, 0, "-", "transforms"]], "cyto_dl.dataframe.readers": [[18, 4, 1, "", "filter_columns"], [18, 4, 1, "", "read_csv"], [18, 4, 1, "", "read_dataframe"], [18, 4, 1, "", "read_h5ad"], [18, 4, 1, "", "read_parquet"]], "cyto_dl.dataframe.transforms": [[20, 0, 0, "-", "filter"], [21, 0, 0, "-", "group_cols"], [22, 0, 0, "-", "misc"], [23, 0, 0, "-", "split"]], "cyto_dl.dataframe.transforms.filter": [[20, 4, 1, "", "filter_columns"], [20, 4, 1, "", "filter_rows"]], "cyto_dl.dataframe.transforms.group_cols": [[21, 1, 1, "", "GroupCols"]], "cyto_dl.dataframe.transforms.misc": [[22, 4, 1, "", "append_class_weights"], [22, 4, 1, "", "append_labels_to_integers"], [22, 4, 1, "", "append_one_hot"], [22, 4, 1, "", "make_random_df"]], "cyto_dl.dataframe.transforms.split": [[23, 4, 1, "", "sample_n_each"], [23, 4, 1, "", "split_dataframe"]], "cyto_dl.datamodules": [[25, 0, 0, "-", "array"], [26, 0, 0, "-", "data_dict"], [27, 0, 0, "-", "dataframe"], [31, 0, 0, "-", "folder"], [32, 0, 0, "-", "multidim_image"], [33, 0, 0, "-", "smartcache"], [34, 0, 0, "-", "torchvision"]], "cyto_dl.datamodules.array": [[25, 4, 1, "", "make_array_dataloader"]], "cyto_dl.datamodules.data_dict": [[26, 4, 1, "", "make_data_dict_dataloader"]], "cyto_dl.datamodules.dataframe": [[28, 0, 0, "-", "dataframe_datamodule"], [29, 0, 0, "-", "grouped_dataframe_datamodule"], [30, 0, 0, "-", "utils"]], "cyto_dl.datamodules.dataframe.dataframe_datamodule": [[28, 1, 1, "", "DataframeDatamodule"]], "cyto_dl.datamodules.dataframe.dataframe_datamodule.DataframeDatamodule": [[28, 2, 1, "", "get_dataloader"], [28, 2, 1, "", "get_dataset"], [28, 2, 1, "", "make_dataloader"], [28, 2, 1, "", "predict_dataloader"], [28, 2, 1, "", "test_dataloader"], [28, 2, 1, "", "train_dataloader"], [28, 2, 1, "", "val_dataloader"]], "cyto_dl.datamodules.dataframe.grouped_dataframe_datamodule": [[29, 1, 1, "", "GroupedDataframeDatamodule"]], "cyto_dl.datamodules.dataframe.grouped_dataframe_datamodule.GroupedDataframeDatamodule": [[29, 2, 1, "", "make_dataloader"]], "cyto_dl.datamodules.dataframe.utils": [[30, 1, 1, "", "AlternatingBatchSampler"], [30, 1, 1, "", "RemoveNaNKeysd"], [30, 4, 1, "", "get_canonical_split_name"], [30, 4, 1, "", "get_dataset"], [30, 4, 1, "", "make_multiple_dataframe_splits"], [30, 4, 1, "", "make_single_dataframe_splits"], [30, 4, 1, "", "parse_transforms"]], "cyto_dl.datamodules.folder": [[31, 4, 1, "", "make_folder_dataloader"]], "cyto_dl.datamodules.multidim_image": [[32, 1, 1, "", "MultiDimImageDataset"], [32, 4, 1, "", "make_multidim_image_dataloader"]], "cyto_dl.datamodules.multidim_image.MultiDimImageDataset": [[32, 2, 1, "", "create_metatensor"], [32, 2, 1, "", "get_per_file_args"], [32, 2, 1, "", "is_batch"]], "cyto_dl.datamodules.smartcache": [[33, 1, 1, "", "SmartcacheDatamodule"]], "cyto_dl.datamodules.smartcache.SmartcacheDatamodule": [[33, 2, 1, "", "get_per_file_args"], [33, 2, 1, "", "make_dataloader"], [33, 2, 1, "", "predict_dataloader"], [33, 2, 1, "", "prepare_data"], [33, 2, 1, "", "setup"], [33, 2, 1, "", "test_dataloader"], [33, 2, 1, "", "train_dataloader"], [33, 2, 1, "", "val_dataloader"]], "cyto_dl.datamodules.torchvision": [[34, 1, 1, "", "CIFAR10"], [34, 1, 1, "", "CIFAR100"], [34, 1, 1, "", "CLEVRClassification"], [34, 1, 1, "", "Caltech101"], [34, 1, 1, "", "Caltech256"], [34, 1, 1, "", "CelebA"], [34, 1, 1, "", "Cityscapes"], [34, 1, 1, "", "CocoCaptions"], [34, 1, 1, "", "CocoDetection"], [34, 1, 1, "", "Country211"], [34, 1, 1, "", "DTD"], [34, 1, 1, "", "DatasetFolder"], [34, 1, 1, "", "EMNIST"], [34, 1, 1, "", "EuroSAT"], [34, 1, 1, "", "FER2013"], [34, 1, 1, "", "FGVCAircraft"], [34, 1, 1, "", "FakeData"], [34, 1, 1, "", "FashionMNIST"], [34, 1, 1, "", "Flickr30k"], [34, 1, 1, "", "Flickr8k"], [34, 1, 1, "", "Flowers102"], [34, 1, 1, "", "Food101"], [34, 1, 1, "", "GTSRB"], [34, 1, 1, "", "HMDB51"], [34, 1, 1, "", "INaturalist"], [34, 1, 1, "", "ImageFolder"], [34, 1, 1, "", "ImageNet"], [34, 1, 1, "", "KMNIST"], [34, 1, 1, "", "Kinetics"], [34, 1, 1, "", "Kitti"], [34, 1, 1, "", "LFWPairs"], [34, 1, 1, "", "LFWPeople"], [34, 1, 1, "", "LSUN"], [34, 1, 1, "", "LSUNClass"], [34, 1, 1, "", "MNIST"], [34, 1, 1, "", "MovingMNIST"], [34, 1, 1, "", "Omniglot"], [34, 1, 1, "", "OxfordIIITPet"], [34, 1, 1, "", "PCAM"], [34, 1, 1, "", "PhotoTour"], [34, 1, 1, "", "Places365"], [34, 1, 1, "", "QMNIST"], [34, 1, 1, "", "RenderedSST2"], [34, 1, 1, "", "SBDataset"], [34, 1, 1, "", "SBU"], [34, 1, 1, "", "SEMEION"], [34, 1, 1, "", "STL10"], [34, 1, 1, "", "SUN397"], [34, 1, 1, "", "SVHN"], [34, 1, 1, "", "StanfordCars"], [34, 1, 1, "", "UCF101"], [34, 1, 1, "", "USPS"], [34, 1, 1, "", "VOCDetection"], [34, 1, 1, "", "VOCSegmentation"], [34, 1, 1, "", "VisionDataset"], [34, 1, 1, "", "WIDERFace"]], "cyto_dl.eval": [[35, 4, 1, "", "main"]], "cyto_dl.image": [[37, 0, 0, "-", "io"], [44, 0, 0, "-", "transforms"]], "cyto_dl.image.io": [[38, 0, 0, "-", "aicsimage_loader"], [39, 0, 0, "-", "monai_bio_reader"], [40, 0, 0, "-", "numpy_reader"], [41, 0, 0, "-", "ome_zarr_reader"], [42, 0, 0, "-", "polygon_loader"], [43, 0, 0, "-", "skimage_reader"]], "cyto_dl.image.io.aicsimage_loader": [[38, 1, 1, "", "AICSImageLoaderd"]], "cyto_dl.image.io.aicsimage_loader.AICSImageLoaderd": [[38, 2, 1, "", "split_args"]], "cyto_dl.image.io.monai_bio_reader": [[39, 1, 1, "", "MonaiBioReader"]], "cyto_dl.image.io.monai_bio_reader.MonaiBioReader": [[39, 2, 1, "", "get_data"], [39, 2, 1, "", "read"], [39, 2, 1, "", "verify_suffix"]], "cyto_dl.image.io.numpy_reader": [[40, 1, 1, "", "ReadNumpyFile"]], "cyto_dl.image.io.ome_zarr_reader": [[41, 1, 1, "", "OmeZarrReader"]], "cyto_dl.image.io.ome_zarr_reader.OmeZarrReader": [[41, 2, 1, "", "get_data"], [41, 2, 1, "", "read"], [41, 2, 1, "", "verify_suffix"]], "cyto_dl.image.io.polygon_loader": [[42, 1, 1, "", "PolygonLoaderd"]], "cyto_dl.image.io.skimage_reader": [[43, 1, 1, "", "SkimageReader"]], "cyto_dl.image.io.skimage_reader.SkimageReader": [[43, 2, 1, "", "get_data"], [43, 2, 1, "", "read"], [43, 2, 1, "", "verify_suffix"]], "cyto_dl.image.transforms": [[45, 0, 0, "-", "add_meta"], [46, 0, 0, "-", "bright_sampler"], [47, 0, 0, "-", "clip"], [48, 0, 0, "-", "contrastadjust"], [49, 0, 0, "-", "generate_jepa_masks"], [50, 0, 0, "-", "merge"], [51, 0, 0, "-", "multiscale_cropper"], [52, 0, 0, "-", "pad"], [53, 0, 0, "-", "project"], [55, 0, 0, "-", "save"], [57, 0, 0, "-", "track_transforms"]], "cyto_dl.image.transforms.add_meta": [[45, 1, 1, "", "AddMeta"]], "cyto_dl.image.transforms.bright_sampler": [[46, 1, 1, "", "BrightSampler"]], "cyto_dl.image.transforms.clip": [[47, 1, 1, "", "Clip"], [47, 1, 1, "", "Clipd"]], "cyto_dl.image.transforms.contrastadjust": [[48, 1, 1, "", "ContrastAdjust"], [48, 1, 1, "", "ContrastAdjustd"]], "cyto_dl.image.transforms.generate_jepa_masks": [[49, 1, 1, "", "JEPAMaskGenerator"]], "cyto_dl.image.transforms.generate_jepa_masks.JEPAMaskGenerator": [[49, 2, 1, "", "remove_excess_pixels"]], "cyto_dl.image.transforms.merge": [[50, 1, 1, "", "Merged"]], "cyto_dl.image.transforms.multiscale_cropper": [[51, 1, 1, "", "RandomMultiScaleCropd"]], "cyto_dl.image.transforms.multiscale_cropper.RandomMultiScaleCropd": [[51, 2, 1, "", "generate_slices"]], "cyto_dl.image.transforms.pad": [[52, 1, 1, "", "PadZd"]], "cyto_dl.image.transforms.project": [[53, 1, 1, "", "MaxProjectd"]], "cyto_dl.image.transforms.save": [[55, 1, 1, "", "Save"], [55, 1, 1, "", "Saved"]], "cyto_dl.image.transforms.track_transforms": [[57, 1, 1, "", "CropResize"], [57, 1, 1, "", "GenerateTrackLabels"], [57, 1, 1, "", "PerChannel"], [57, 1, 1, "", "SplitTrackd"]], "cyto_dl.loggers": [[59, 0, 0, "-", "mlflow"]], "cyto_dl.loggers.mlflow": [[59, 1, 1, "", "MLFlowLogger"]], "cyto_dl.loggers.mlflow.MLFlowLogger": [[59, 2, 1, "", "after_save_checkpoint"], [59, 2, 1, "", "log_hyperparams"], [59, 2, 1, "", "log_metrics"]], "cyto_dl.models": [[61, 0, 0, "-", "base_model"], [62, 0, 0, "-", "basic_model"], [63, 0, 0, "-", "classification"], [66, 0, 0, "-", "contrastive"], [72, 0, 0, "-", "im2im"], [82, 0, 0, "-", "jepa"], [86, 0, 0, "-", "utils"]], "cyto_dl.models.base_model": [[61, 1, 1, "", "BaseModel"], [61, 1, 1, "", "BaseModelMeta"]], "cyto_dl.models.base_model.BaseModel": [[61, 2, 1, "", "compute_metrics"], [61, 2, 1, "", "configure_optimizers"], [61, 2, 1, "", "forward"], [61, 2, 1, "", "model_step"], [61, 2, 1, "", "on_train_start"], [61, 2, 1, "", "parse_batch"], [61, 2, 1, "", "predict_step"], [61, 2, 1, "", "test_step"], [61, 2, 1, "", "training_step"], [61, 2, 1, "", "validation_step"]], "cyto_dl.models.basic_model": [[62, 1, 1, "", "BasicModel"]], "cyto_dl.models.basic_model.BasicModel": [[62, 2, 1, "", "forward"], [62, 2, 1, "", "model_step"]], "cyto_dl.models.classification": [[64, 0, 0, "-", "classification"], [65, 0, 0, "-", "timepoint_classification"]], "cyto_dl.models.classification.classification": [[64, 1, 1, "", "Classification"]], "cyto_dl.models.classification.classification.Classification": [[64, 2, 1, "", "forward"], [64, 2, 1, "", "model_step"], [64, 2, 1, "", "predict_step"], [64, 2, 1, "", "save_images"], [64, 2, 1, "", "should_save_image"]], "cyto_dl.models.classification.timepoint_classification": [[65, 1, 1, "", "TimepointClassification"]], "cyto_dl.models.classification.timepoint_classification.TimepointClassification": [[65, 2, 1, "", "predict_step"]], "cyto_dl.models.contrastive": [[67, 0, 0, "-", "contrastive"]], "cyto_dl.models.contrastive.contrastive": [[67, 1, 1, "", "Contrastive"]], "cyto_dl.models.contrastive.contrastive.Contrastive": [[67, 2, 1, "", "forward"], [67, 2, 1, "", "model_step"], [67, 2, 1, "", "plot_classes"], [67, 2, 1, "", "plot_neighbors"], [67, 2, 1, "", "predict_step"]], "cyto_dl.models.im2im": [[73, 0, 0, "-", "gan"], [74, 0, 0, "-", "multi_task"], [75, 0, 0, "-", "utils"]], "cyto_dl.models.im2im.gan": [[73, 1, 1, "", "GAN"]], "cyto_dl.models.im2im.gan.GAN": [[73, 2, 1, "", "configure_optimizers"], [73, 2, 1, "", "model_step"], [73, 2, 1, "", "predict_step"]], "cyto_dl.models.im2im.multi_task": [[74, 1, 1, "", "MultiTaskIm2Im"]], "cyto_dl.models.im2im.multi_task.MultiTaskIm2Im": [[74, 2, 1, "", "configure_optimizers"], [74, 2, 1, "", "forward"], [74, 2, 1, "", "get_n_postprocess_image"], [74, 2, 1, "", "get_per_head"], [74, 2, 1, "", "model_step"], [74, 2, 1, "", "predict_step"], [74, 2, 1, "", "run_forward"], [74, 2, 1, "", "test_step"], [74, 2, 1, "", "training_step"], [74, 2, 1, "", "validation_step"]], "cyto_dl.models.im2im.utils": [[76, 0, 0, "-", "instance_seg"], [77, 0, 0, "-", "noise_annealer"], [78, 0, 0, "-", "postprocessing"]], "cyto_dl.models.im2im.utils.instance_seg": [[76, 1, 1, "", "InstanceSegCluster"], [76, 1, 1, "", "InstanceSegLoss"], [76, 1, 1, "", "InstanceSegPreprocessd"], [76, 1, 1, "", "InstanceSegRandFlipd"], [76, 4, 1, "", "pad_slice"]], "cyto_dl.models.im2im.utils.instance_seg.InstanceSegCluster": [[76, 2, 1, "", "cluster_object"], [76, 2, 1, "", "kd_clustering"], [76, 2, 1, "", "remove_small_skeletons"]], "cyto_dl.models.im2im.utils.instance_seg.InstanceSegPreprocessd": [[76, 2, 1, "", "embed_from_skel"], [76, 2, 1, "", "keep_largest_cc"], [76, 2, 1, "", "label_2d"], [76, 2, 1, "", "shrink"], [76, 2, 1, "", "skeleton_tall"], [76, 2, 1, "", "smooth_embedding"], [76, 2, 1, "", "topology_preserving_thinning"]], "cyto_dl.models.im2im.utils.noise_annealer": [[77, 1, 1, "", "NoiseAnnealer"]], "cyto_dl.models.im2im.utils.noise_annealer.NoiseAnnealer": [[77, 2, 1, "", "update_noise"]], "cyto_dl.models.im2im.utils.postprocessing": [[79, 0, 0, "-", "act_thresh_label"], [80, 0, 0, "-", "auto_thresh"], [81, 0, 0, "-", "dict_to_im"]], "cyto_dl.models.im2im.utils.postprocessing.act_thresh_label": [[79, 1, 1, "", "ActThreshLabel"]], "cyto_dl.models.im2im.utils.postprocessing.auto_thresh": [[80, 1, 1, "", "AutoThreshold"]], "cyto_dl.models.im2im.utils.postprocessing.dict_to_im": [[81, 1, 1, "", "DictToIm"], [81, 4, 1, "", "detach"]], "cyto_dl.models.jepa": [[83, 0, 0, "-", "ijepa"], [84, 0, 0, "-", "iwm"], [85, 0, 0, "-", "jepa_base"]], "cyto_dl.models.jepa.ijepa": [[83, 1, 1, "", "IJEPA"]], "cyto_dl.models.jepa.ijepa.IJEPA": [[83, 2, 1, "", "model_step"]], "cyto_dl.models.jepa.iwm": [[84, 1, 1, "", "IWM"]], "cyto_dl.models.jepa.iwm.IWM": [[84, 2, 1, "", "extract_embeddings"], [84, 2, 1, "", "get_predict_masks"], [84, 2, 1, "", "model_step"], [84, 2, 1, "", "predict_step"], [84, 2, 1, "", "test_step"]], "cyto_dl.models.jepa.jepa_base": [[85, 1, 1, "", "JEPABase"]], "cyto_dl.models.jepa.jepa_base.JEPABase": [[85, 2, 1, "", "configure_optimizers"], [85, 2, 1, "", "forward"], [85, 2, 1, "", "get_context_embeddings"], [85, 2, 1, "", "get_mask"], [85, 2, 1, "", "get_target_embeddings"], [85, 2, 1, "", "model_step"], [85, 2, 1, "", "predict_step"], [85, 2, 1, "", "remove_first_dim"], [85, 2, 1, "", "update_teacher"]], "cyto_dl.models.utils": [[87, 0, 0, "-", "mlflow"], [88, 0, 0, "-", "utils"]], "cyto_dl.models.utils.mlflow": [[87, 4, 1, "", "get_config"], [87, 4, 1, "", "load_model_from_checkpoint"]], "cyto_dl.models.utils.utils": [[88, 4, 1, "", "find_indices"]], "cyto_dl.models.vae": [[90, 0, 0, "-", "base_vae"], [99, 0, 0, "-", "priors"]], "cyto_dl.models.vae.base_vae": [[90, 1, 1, "", "BaseVAE"]], "cyto_dl.models.vae.base_vae.BaseVAE": [[90, 2, 1, "", "calculate_elbo"], [90, 2, 1, "", "calculate_rcl"], [90, 2, 1, "", "calculate_rcl_dict"], [90, 2, 1, "", "decode"], [90, 2, 1, "", "encode"], [90, 2, 1, "", "forward"], [90, 2, 1, "", "model_step"], [90, 2, 1, "", "sample_z"]], "cyto_dl.models.vae.priors": [[100, 0, 0, "-", "abstract_prior"], [101, 0, 0, "-", "gaussian"], [102, 0, 0, "-", "identity_prior"], [103, 0, 0, "-", "joint_prior"]], "cyto_dl.models.vae.priors.abstract_prior": [[100, 1, 1, "", "Prior"]], "cyto_dl.models.vae.priors.abstract_prior.Prior": [[100, 2, 1, "", "forward"], [100, 2, 1, "", "kl_divergence"], [100, 5, 1, "", "param_size"], [100, 2, 1, "", "sample"]], "cyto_dl.models.vae.priors.gaussian": [[101, 1, 1, "", "DiagonalGaussianPrior"], [101, 1, 1, "", "IsotropicGaussianPrior"], [101, 4, 1, "", "compute_tc_penalty"]], "cyto_dl.models.vae.priors.gaussian.DiagonalGaussianPrior": [[101, 2, 1, "", "forward"], [101, 2, 1, "", "kl_divergence"], [101, 5, 1, "", "param_size"]], "cyto_dl.models.vae.priors.gaussian.IsotropicGaussianPrior": [[101, 2, 1, "", "forward"], [101, 2, 1, "", "kl_divergence"], [101, 5, 1, "", "param_size"], [101, 2, 1, "", "sample"]], "cyto_dl.models.vae.priors.identity_prior": [[102, 1, 1, "", "IdentityPrior"]], "cyto_dl.models.vae.priors.identity_prior.IdentityPrior": [[102, 2, 1, "", "forward"], [102, 5, 1, "", "param_size"]], "cyto_dl.models.vae.priors.joint_prior": [[103, 1, 1, "", "JointPrior"]], "cyto_dl.models.vae.priors.joint_prior.JointPrior": [[103, 2, 1, "", "forward"], [103, 2, 1, "", "kl_divergence"], [103, 5, 1, "", "param_size"], [103, 2, 1, "", "sample"]], "cyto_dl.nn": [[106, 0, 0, "-", "discriminators"], [109, 0, 0, "-", "head"], [117, 0, 0, "-", "hr_skip"], [118, 0, 0, "-", "losses"], [132, 0, 0, "-", "mlp"], [138, 0, 0, "-", "res_unit"], [139, 0, 0, "-", "spatial_transformer"], [140, 0, 0, "-", "track_sequence_predictor"], [141, 0, 0, "-", "vits"]], "cyto_dl.nn.discriminators": [[107, 0, 0, "-", "multi_scale_discriminator"], [108, 0, 0, "-", "n_layer_discriminator"]], "cyto_dl.nn.discriminators.multi_scale_discriminator": [[107, 1, 1, "", "MultiScaleDiscriminator"]], "cyto_dl.nn.discriminators.multi_scale_discriminator.MultiScaleDiscriminator": [[107, 2, 1, "", "forward"]], "cyto_dl.nn.discriminators.n_layer_discriminator": [[108, 1, 1, "", "NLayerDiscriminator"]], "cyto_dl.nn.discriminators.n_layer_discriminator.NLayerDiscriminator": [[108, 2, 1, "", "forward"]], "cyto_dl.nn.head": [[110, 0, 0, "-", "base_head"], [111, 0, 0, "-", "gan_head"], [112, 0, 0, "-", "gan_head_superres"], [113, 0, 0, "-", "mae_head"], [114, 0, 0, "-", "mask_head"], [115, 0, 0, "-", "res_blocks_head"], [116, 0, 0, "-", "vic_reg"]], "cyto_dl.nn.head.base_head": [[110, 1, 1, "", "BaseHead"]], "cyto_dl.nn.head.base_head.BaseHead": [[110, 2, 1, "", "forward"], [110, 2, 1, "", "generate_io_map"], [110, 2, 1, "", "run_head"], [110, 2, 1, "", "update_params"]], "cyto_dl.nn.head.gan_head": [[111, 1, 1, "", "GANHead"]], "cyto_dl.nn.head.gan_head.GANHead": [[111, 2, 1, "", "forward"], [111, 2, 1, "", "run_head"]], "cyto_dl.nn.head.gan_head_superres": [[112, 1, 1, "", "GANHead_resize"]], "cyto_dl.nn.head.gan_head_superres.GANHead_resize": [[112, 2, 1, "", "forward"]], "cyto_dl.nn.head.mae_head": [[113, 1, 1, "", "MAEHead"]], "cyto_dl.nn.head.mae_head.MAEHead": [[113, 2, 1, "", "run_head"]], "cyto_dl.nn.head.mask_head": [[114, 1, 1, "", "MaskHead"]], "cyto_dl.nn.head.mask_head.MaskHead": [[114, 2, 1, "", "run_head"]], "cyto_dl.nn.head.res_blocks_head": [[115, 1, 1, "", "ResBlocksHead"]], "cyto_dl.nn.head.res_blocks_head.ResBlocksHead": [[115, 2, 1, "", "forward"]], "cyto_dl.nn.head.vic_reg": [[116, 1, 1, "", "VICRegHead"]], "cyto_dl.nn.head.vic_reg.VICRegHead": [[116, 2, 1, "", "run_head"]], "cyto_dl.nn.hr_skip": [[117, 1, 1, "", "HRSkip"]], "cyto_dl.nn.hr_skip.HRSkip": [[117, 2, 1, "", "forward"]], "cyto_dl.nn.losses": [[119, 0, 0, "-", "adversarial_loss"], [120, 0, 0, "-", "chamfer_loss"], [121, 0, 0, "-", "continuous_bernoulli"], [122, 0, 0, "-", "cosine_loss"], [123, 0, 0, "-", "gan_loss"], [124, 0, 0, "-", "gaussian_nll_loss"], [126, 0, 0, "-", "loss_wrapper"], [128, 0, 0, "-", "threshold_loss"], [129, 0, 0, "-", "vic_reg"], [130, 0, 0, "-", "weibull"], [131, 0, 0, "-", "weighted_mse_loss"]], "cyto_dl.nn.losses.adversarial_loss": [[119, 1, 1, "", "AdversarialLoss"]], "cyto_dl.nn.losses.adversarial_loss.AdversarialLoss": [[119, 2, 1, "", "forward"]], "cyto_dl.nn.losses.chamfer_loss": [[120, 1, 1, "", "ChamferLoss"]], "cyto_dl.nn.losses.chamfer_loss.ChamferLoss": [[120, 2, 1, "", "batch_pairwise_dist"], [120, 2, 1, "", "batch_pairwise_dist2"], [120, 2, 1, "", "forward"]], "cyto_dl.nn.losses.continuous_bernoulli": [[121, 1, 1, "", "CBLogLoss"]], "cyto_dl.nn.losses.continuous_bernoulli.CBLogLoss": [[121, 2, 1, "", "forward"]], "cyto_dl.nn.losses.cosine_loss": [[122, 1, 1, "", "CosineLoss"]], "cyto_dl.nn.losses.cosine_loss.CosineLoss": [[122, 2, 1, "", "forward"]], "cyto_dl.nn.losses.gan_loss": [[123, 1, 1, "", "GANLoss"], [123, 1, 1, "", "Pix2PixHD"]], "cyto_dl.nn.losses.gan_loss.GANLoss": [[123, 2, 1, "", "get_target_tensor"]], "cyto_dl.nn.losses.gan_loss.Pix2PixHD": [[123, 2, 1, "", "get_feature_matching_loss"], [123, 2, 1, "", "get_gan_loss"]], "cyto_dl.nn.losses.gaussian_nll_loss": [[124, 1, 1, "", "GaussianNLLLoss"]], "cyto_dl.nn.losses.gaussian_nll_loss.GaussianNLLLoss": [[124, 2, 1, "", "forward"]], "cyto_dl.nn.losses.loss_wrapper": [[126, 1, 1, "", "CMAP_loss"], [126, 1, 1, "", "LossWrapper"]], "cyto_dl.nn.losses.threshold_loss": [[128, 1, 1, "", "ThresholdLoss"]], "cyto_dl.nn.losses.vic_reg": [[129, 1, 1, "", "VICRegLoss"], [129, 4, 1, "", "off_diagonal"]], "cyto_dl.nn.losses.vic_reg.VICRegLoss": [[129, 2, 1, "", "forward"]], "cyto_dl.nn.losses.weibull": [[130, 1, 1, "", "WeibullLogLoss"], [130, 4, 1, "", "weibull_log_probs"]], "cyto_dl.nn.losses.weibull.WeibullLogLoss": [[130, 2, 1, "", "forward"]], "cyto_dl.nn.losses.weighted_mse_loss": [[131, 1, 1, "", "WeightedMSELoss"]], "cyto_dl.nn.losses.weighted_mse_loss.WeightedMSELoss": [[131, 2, 1, "", "forward"]], "cyto_dl.nn.mlp": [[132, 1, 1, "", "MLP"]], "cyto_dl.nn.mlp.MLP": [[132, 2, 1, "", "forward"]], "cyto_dl.nn.res_unit": [[138, 1, 1, "", "ResidualUnit"]], "cyto_dl.nn.res_unit.ResidualUnit": [[138, 2, 1, "", "forward"]], "cyto_dl.nn.spatial_transformer": [[139, 1, 1, "", "ConvPoolReLU"], [139, 1, 1, "", "STN"]], "cyto_dl.nn.spatial_transformer.ConvPoolReLU": [[139, 2, 1, "", "forward"]], "cyto_dl.nn.spatial_transformer.STN": [[139, 2, 1, "", "forward"], [139, 2, 1, "", "toggle"]], "cyto_dl.nn.track_sequence_predictor": [[140, 1, 1, "", "TrackClassifier"], [140, 4, 1, "", "positionalencoding1d"]], "cyto_dl.nn.track_sequence_predictor.TrackClassifier": [[140, 2, 1, "", "forward"]], "cyto_dl.nn.vits": [[142, 0, 0, "-", "blocks"], [146, 0, 0, "-", "cross_mae"], [147, 0, 0, "-", "jepa"], [148, 0, 0, "-", "mae"], [149, 0, 0, "-", "seg"], [150, 0, 0, "-", "utils"]], "cyto_dl.nn.vits.blocks": [[143, 0, 0, "-", "cross_attention"], [144, 0, 0, "-", "intermediate_weigher"], [145, 0, 0, "-", "patchify"]], "cyto_dl.nn.vits.blocks.cross_attention": [[143, 1, 1, "", "CrossAttention"], [143, 1, 1, "", "CrossAttentionBlock"], [143, 1, 1, "", "CrossSelfBlock"], [143, 1, 1, "", "Mlp"]], "cyto_dl.nn.vits.blocks.cross_attention.CrossAttention": [[143, 2, 1, "", "forward"]], "cyto_dl.nn.vits.blocks.cross_attention.CrossAttentionBlock": [[143, 2, 1, "", "forward"]], "cyto_dl.nn.vits.blocks.cross_attention.CrossSelfBlock": [[143, 2, 1, "", "forward"]], "cyto_dl.nn.vits.blocks.cross_attention.Mlp": [[143, 2, 1, "", "forward"]], "cyto_dl.nn.vits.blocks.intermediate_weigher": [[144, 1, 1, "", "IntermediateWeigher"]], "cyto_dl.nn.vits.blocks.intermediate_weigher.IntermediateWeigher": [[144, 2, 1, "", "forward"]], "cyto_dl.nn.vits.blocks.patchify": [[145, 1, 1, "", "Patchify"], [145, 4, 1, "", "random_indexes"]], "cyto_dl.nn.vits.blocks.patchify.Patchify": [[145, 2, 1, "", "forward"], [145, 2, 1, "", "get_mask"]], "cyto_dl.nn.vits.cross_mae": [[146, 1, 1, "", "CrossMAE_Decoder"]], "cyto_dl.nn.vits.cross_mae.CrossMAE_Decoder": [[146, 2, 1, "", "forward"], [146, 2, 1, "", "init_weight"]], "cyto_dl.nn.vits.jepa": [[147, 1, 1, "", "IWMPredictor"], [147, 1, 1, "", "JEPAEncoder"], [147, 1, 1, "", "JEPAPredictor"]], "cyto_dl.nn.vits.jepa.IWMPredictor": [[147, 2, 1, "", "forward"]], "cyto_dl.nn.vits.jepa.JEPAEncoder": [[147, 2, 1, "", "forward"]], "cyto_dl.nn.vits.jepa.JEPAPredictor": [[147, 2, 1, "", "forward"], [147, 2, 1, "", "init_weight"], [147, 2, 1, "", "predict_target_features"]], "cyto_dl.nn.vits.mae": [[148, 1, 1, "", "MAE_Decoder"], [148, 1, 1, "", "MAE_Encoder"], [148, 1, 1, "", "MAE_ViT"]], "cyto_dl.nn.vits.mae.MAE_Decoder": [[148, 2, 1, "", "forward"], [148, 2, 1, "", "init_weight"]], "cyto_dl.nn.vits.mae.MAE_Encoder": [[148, 2, 1, "", "forward"], [148, 2, 1, "", "init_weight"]], "cyto_dl.nn.vits.mae.MAE_ViT": [[148, 2, 1, "", "forward"]], "cyto_dl.nn.vits.seg": [[149, 1, 1, "", "EncodedSkip"], [149, 1, 1, "", "Seg_ViT"], [149, 1, 1, "", "SuperresDecoder"]], "cyto_dl.nn.vits.seg.EncodedSkip": [[149, 2, 1, "", "forward"]], "cyto_dl.nn.vits.seg.Seg_ViT": [[149, 2, 1, "", "forward"]], "cyto_dl.nn.vits.seg.SuperresDecoder": [[149, 2, 1, "", "forward"]], "cyto_dl.nn.vits.utils": [[150, 4, 1, "", "get_positional_embedding"], [150, 4, 1, "", "take_indexes"]], "cyto_dl.train": [[154, 4, 1, "", "main"]], "cyto_dl.utils": [[156, 0, 0, "-", "array"], [157, 0, 0, "-", "checkpoint"], [158, 0, 0, "-", "config"], [159, 0, 0, "-", "download_test_data"], [160, 0, 0, "-", "dummy_dataset"], [161, 0, 0, "-", "pylogger"], [162, 0, 0, "-", "rich_utils"], [163, 0, 0, "-", "rotation"], [167, 0, 0, "-", "template_utils"]], "cyto_dl.utils.array": [[156, 4, 1, "", "create_dataloader"], [156, 4, 1, "", "extract_array_predictions"]], "cyto_dl.utils.checkpoint": [[157, 4, 1, "", "load_checkpoint"]], "cyto_dl.utils.config": [[158, 4, 1, "", "is_config"], [158, 4, 1, "", "kv_to_dict"], [158, 4, 1, "", "remove_aux_key"]], "cyto_dl.utils.download_test_data": [[159, 4, 1, "", "delete_test_data"], [159, 4, 1, "", "download_test_data"], [159, 4, 1, "", "parse_s3_path"], [159, 4, 1, "", "setup_paths"]], "cyto_dl.utils.dummy_dataset": [[160, 1, 1, "", "DummyDatamodule"], [160, 1, 1, "", "DummyDataset"]], "cyto_dl.utils.dummy_dataset.DummyDatamodule": [[160, 2, 1, "", "get_dataloader"], [160, 2, 1, "", "predict_dataloader"], [160, 2, 1, "", "test_dataloader"], [160, 2, 1, "", "train_dataloader"], [160, 2, 1, "", "val_dataloader"]], "cyto_dl.utils.dummy_dataset.DummyDataset": [[160, 2, 1, "", "generate_img"]], "cyto_dl.utils.pylogger": [[161, 4, 1, "", "get_pylogger"]], "cyto_dl.utils.rich_utils": [[162, 4, 1, "", "enforce_tags"], [162, 4, 1, "", "print_config_tree"]], "cyto_dl.utils.rotation": [[163, 1, 1, "", "RotationModule"], [163, 4, 1, "", "get_rotation_matrix_so2"], [163, 4, 1, "", "get_rotation_matrix_so3"]], "cyto_dl.utils.rotation.RotationModule": [[163, 2, 1, "", "compute_rotation_matrix"]], "cyto_dl.utils.template_utils": [[167, 4, 1, "", "close_loggers"], [167, 4, 1, "", "extras"], [167, 4, 1, "", "get_metric_value"], [167, 4, 1, "", "instantiate_callbacks"], [167, 4, 1, "", "instantiate_loggers"], [167, 4, 1, "", "log_hyperparameters"], [167, 4, 1, "", "save_file"], [167, 4, 1, "", "task_wrapper"]]}, "objtypes": {"0": "py:module", "1": "py:class", "2": "py:method", "3": "py:attribute", "4": "py:function", "5": "py:property"}, "objnames": {"0": ["py", "module", "Python module"], "1": ["py", "class", "Python class"], "2": ["py", "method", "Python method"], "3": ["py", "attribute", "Python attribute"], "4": ["py", "function", "Python function"], "5": ["py", "property", "Python property"]}, "titleterms": {"us": 0, "api": [0, 2, 3, 4, 5, 6, 7], "load": 0, "config": [0, 158, 172], "train": [0, 154, 172], "predict": [0, 172], "util": [0, 30, 75, 76, 77, 78, 79, 80, 81, 86, 87, 88, 150, 155, 156, 157, 158, 159, 160, 161, 162, 163, 164, 165, 166, 167], "cyto_dl": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153, 154, 155, 156, 157, 158, 159, 160, 161, 162, 163, 164, 165, 166, 167, 170], "packag": [1, 2, 3, 8, 17, 19, 24, 27, 36, 37, 44, 58, 60, 63, 66, 68, 72, 75, 78, 82, 86, 89, 95, 99, 105, 106, 109, 118, 133, 141, 142, 151, 152, 155, 164], "subpackag": [1, 2, 17, 24, 36, 60, 72, 75, 89, 105, 141, 151, 155], "submodul": [1, 2, 3, 8, 17, 19, 24, 27, 37, 44, 58, 60, 63, 66, 68, 72, 75, 78, 82, 86, 89, 95, 99, 105, 106, 109, 118, 133, 141, 142, 152, 155, 164], "cyto_dl_model": [3, 4, 5], "cyto_dl_base_model": 4, "modul": [4, 5, 6, 7, 9, 10, 11, 12, 13, 14, 15, 16, 18, 20, 21, 22, 23, 25, 26, 28, 29, 30, 31, 32, 33, 34, 35, 38, 39, 40, 41, 42, 43, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 59, 61, 62, 64, 65, 67, 69, 70, 71, 73, 74, 76, 77, 79, 80, 81, 83, 84, 85, 87, 88, 90, 91, 92, 93, 94, 96, 97, 98, 100, 101, 102, 103, 104, 107, 108, 110, 111, 112, 113, 114, 115, 116, 117, 119, 120, 121, 122, 123, 124, 125, 126, 127, 128, 129, 130, 131, 132, 134, 135, 136, 137, 138, 139, 140, 143, 144, 145, 146, 147, 148, 149, 150, 153, 154, 156, 157, 158, 159, 160, 161, 162, 163, 165, 166, 167], "segmentation_plugin_model": 5, "data": [6, 168, 171], "model": [7, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 168, 171], "callback": [8, 9, 10, 11, 12, 13, 14, 15], "callback_util": 9, "csv_saver": 10, "image_sav": 11, "latent_walk": 12, "layer_freez": 13, "model_util": 14, "outlier_detect": 15, "compil": 16, "datafram": [17, 18, 19, 20, 21, 22, 23, 27, 28, 29, 30], "reader": 18, "transform": [19, 20, 21, 22, 23, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57], "filter": 20, "group_col": 21, "misc": 22, "split": 23, "datamodul": [24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34], "arrai": [25, 156], "data_dict": 26, "dataframe_datamodul": 28, "grouped_dataframe_datamodul": 29, "folder": 31, "multidim_imag": 32, "smartcach": 33, "torchvis": 34, "eval": 35, "imag": [36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57], "io": [37, 38, 39, 40, 41, 42, 43, 152, 153], "aicsimage_load": 38, "monai_bio_read": 39, "numpy_read": 40, "ome_zarr_read": 41, "polygon_load": 42, "skimage_read": 43, "add_meta": 45, "bright_sampl": 46, "clip": 47, "contrastadjust": 48, "generate_jepa_mask": 49, "merg": 50, "multiscale_cropp": 51, "pad": 52, "project": 53, "rotation_mask_transform": 54, "save": 55, "so2_random_rot": 56, "track_transform": 57, "logger": [58, 59], "mlflow": [59, 87], "base_model": 61, "basic_model": 62, "classif": [63, 64, 65], "timepoint_classif": 65, "contrast": [66, 67], "handler": [68, 69, 70, 71], "base_handl": 69, "image_handl": 70, "load_image_patch": 71, "im2im": [72, 73, 74, 75, 76, 77, 78, 79, 80, 81], "gan": 73, "multi_task": 74, "instance_seg": 76, "noise_anneal": 77, "postprocess": [78, 79, 80, 81], "act_thresh_label": 79, "auto_thresh": 80, "dict_to_im": 81, "jepa": [82, 83, 84, 85, 147], "ijepa": 83, "iwm": 84, "jepa_bas": 85, "vae": [89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104], "base_va": 90, "image_canon_va": 91, "image_encod": 92, "image_va": 93, "latent_loss_va": 94, "o2_spharm_va": [95, 96, 97], "o2_spharm_encod": 96, "point_cloud_va": 98, "prior": [99, 100, 101, 102, 103], "abstract_prior": 100, "gaussian": 101, "identity_prior": 102, "joint_prior": 103, "tabular_va": 104, "nn": [105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150], "discrimin": [106, 107, 108], "multi_scale_discrimin": 107, "n_layer_discrimin": 108, "head": [109, 110, 111, 112, 113, 114, 115, 116], "base_head": 110, "gan_head": 111, "gan_head_superr": 112, "mae_head": 113, "mask_head": 114, "res_blocks_head": 115, "vic_reg": [116, 129], "hr_skip": 117, "loss": [118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 128, 129, 130, 131], "adversarial_loss": 119, "chamfer_loss": 120, "continuous_bernoulli": 121, "cosine_loss": 122, "gan_loss": 123, "gaussian_nll_loss": 124, "geomloss": 125, "loss_wrapp": 126, "spharm_loss": 127, "threshold_loss": 128, "weibul": 130, "weighted_mse_loss": 131, "mlp": 132, "point_cloud": [133, 134, 135, 136, 137, 151, 152, 153], "dgcnn": 134, "folding_net": 135, "graph_funct": 136, "vnn": 137, "res_unit": 138, "spatial_transform": 139, "track_sequence_predictor": 140, "vit": [141, 142, 143, 144, 145, 146, 147, 148, 149, 150], "block": [142, 143, 144, 145], "cross_attent": 143, "intermediate_weigh": 144, "patchifi": 145, "cross_ma": 146, "mae": 148, "seg": 149, "read_pcloud": 153, "checkpoint": 157, "download_test_data": 159, "dummy_dataset": 160, "pylogg": 161, "rich_util": 162, "rotat": [163, 166], "spharm": [164, 165, 166], "reconstruct": 165, "template_util": 167, "overview": [168, 171], "experi": [168, 171], "instal": 169, "modifi": 172, "exampl": 172, "test": 172}, "envversion": {"sphinx.domains.c": 2, "sphinx.domains.changeset": 1, "sphinx.domains.citation": 1, "sphinx.domains.cpp": 8, "sphinx.domains.index": 1, "sphinx.domains.javascript": 2, "sphinx.domains.math": 2, "sphinx.domains.python": 3, "sphinx.domains.rst": 2, "sphinx.domains.std": 2, "sphinx.ext.viewcode": 1, "sphinx": 57}, "alltitles": {"Using the API": [[0, "using-the-api"]], "Loading Configs": [[0, "loading-configs"]], "Training and Prediction": [[0, "training-and-prediction"]], "Utils": [[0, "utils"]], "cyto_dl package": [[1, "module-cyto_dl"]], "Subpackages": [[1, "subpackages"], [2, "subpackages"], [17, "subpackages"], [24, "subpackages"], [36, "subpackages"], [60, "subpackages"], [72, "subpackages"], [75, "subpackages"], [89, "subpackages"], [105, "subpackages"], [141, "subpackages"], [151, "subpackages"], [155, "subpackages"]], "Submodules": [[1, "submodules"], [2, "submodules"], [3, "submodules"], [8, "submodules"], [17, "submodules"], [19, "submodules"], [24, "submodules"], [27, "submodules"], [37, "submodules"], [44, "submodules"], [58, "submodules"], [60, "submodules"], [63, "submodules"], [66, "submodules"], [68, "submodules"], [72, "submodules"], [75, "submodules"], [78, "submodules"], [82, "submodules"], [86, "submodules"], [89, "submodules"], [95, "submodules"], [99, "submodules"], [105, "submodules"], [106, "submodules"], [109, "submodules"], [118, "submodules"], [133, "submodules"], [141, "submodules"], [142, "submodules"], [152, "submodules"], [155, "submodules"], [164, "submodules"]], "cyto_dl.api package": [[2, "module-cyto_dl.api"]], "cyto_dl.api.cyto_dl_model package": [[3, "module-cyto_dl.api.cyto_dl_model"]], "cyto_dl.api.cyto_dl_model.cyto_dl_base_model module": [[4, "module-cyto_dl.api.cyto_dl_model.cyto_dl_base_model"]], "cyto_dl.api.cyto_dl_model.segmentation_plugin_model module": [[5, "module-cyto_dl.api.cyto_dl_model.segmentation_plugin_model"]], "cyto_dl.api.data module": [[6, "module-cyto_dl.api.data"]], "cyto_dl.api.model module": [[7, "module-cyto_dl.api.model"]], "cyto_dl.callbacks package": [[8, "module-cyto_dl.callbacks"]], "cyto_dl.callbacks.callback_utils module": [[9, "cyto-dl-callbacks-callback-utils-module"]], "cyto_dl.callbacks.csv_saver module": [[10, "module-cyto_dl.callbacks.csv_saver"]], "cyto_dl.callbacks.image_saver module": [[11, "module-cyto_dl.callbacks.image_saver"]], "cyto_dl.callbacks.latent_walk module": [[12, "cyto-dl-callbacks-latent-walk-module"]], "cyto_dl.callbacks.layer_freeze module": [[13, "module-cyto_dl.callbacks.layer_freeze"]], "cyto_dl.callbacks.model_utils module": [[14, "module-cyto_dl.callbacks.model_utils"]], "cyto_dl.callbacks.outlier_detection module": [[15, "module-cyto_dl.callbacks.outlier_detection"]], "cyto_dl.compile module": [[16, "cyto-dl-compile-module"]], "cyto_dl.dataframe package": [[17, "module-cyto_dl.dataframe"]], "cyto_dl.dataframe.readers module": [[18, "module-cyto_dl.dataframe.readers"]], "cyto_dl.dataframe.transforms package": [[19, "module-cyto_dl.dataframe.transforms"]], "cyto_dl.dataframe.transforms.filter module": [[20, "module-cyto_dl.dataframe.transforms.filter"]], "cyto_dl.dataframe.transforms.group_cols module": [[21, "module-cyto_dl.dataframe.transforms.group_cols"]], "cyto_dl.dataframe.transforms.misc module": [[22, "module-cyto_dl.dataframe.transforms.misc"]], "cyto_dl.dataframe.transforms.split module": [[23, "module-cyto_dl.dataframe.transforms.split"]], "cyto_dl.datamodules package": [[24, "module-cyto_dl.datamodules"]], "cyto_dl.datamodules.array module": [[25, "module-cyto_dl.datamodules.array"]], "cyto_dl.datamodules.data_dict module": [[26, "module-cyto_dl.datamodules.data_dict"]], "cyto_dl.datamodules.dataframe package": [[27, "module-cyto_dl.datamodules.dataframe"]], "cyto_dl.datamodules.dataframe.dataframe_datamodule module": [[28, "module-cyto_dl.datamodules.dataframe.dataframe_datamodule"]], "cyto_dl.datamodules.dataframe.grouped_dataframe_datamodule module": [[29, "module-cyto_dl.datamodules.dataframe.grouped_dataframe_datamodule"]], "cyto_dl.datamodules.dataframe.utils module": [[30, "module-cyto_dl.datamodules.dataframe.utils"]], "cyto_dl.datamodules.folder module": [[31, "module-cyto_dl.datamodules.folder"]], "cyto_dl.datamodules.multidim_image module": [[32, "module-cyto_dl.datamodules.multidim_image"]], "cyto_dl.datamodules.smartcache module": [[33, "module-cyto_dl.datamodules.smartcache"]], "cyto_dl.datamodules.torchvision module": [[34, "module-cyto_dl.datamodules.torchvision"]], "cyto_dl.eval module": [[35, "module-cyto_dl.eval"]], "cyto_dl.image package": [[36, "module-cyto_dl.image"]], "cyto_dl.image.io package": [[37, "module-cyto_dl.image.io"]], "cyto_dl.image.io.aicsimage_loader module": [[38, "module-cyto_dl.image.io.aicsimage_loader"]], "cyto_dl.image.io.monai_bio_reader module": [[39, "module-cyto_dl.image.io.monai_bio_reader"]], "cyto_dl.image.io.numpy_reader module": [[40, "module-cyto_dl.image.io.numpy_reader"]], "cyto_dl.image.io.ome_zarr_reader module": [[41, "module-cyto_dl.image.io.ome_zarr_reader"]], "cyto_dl.image.io.polygon_loader module": [[42, "module-cyto_dl.image.io.polygon_loader"]], "cyto_dl.image.io.skimage_reader module": [[43, "module-cyto_dl.image.io.skimage_reader"]], "cyto_dl.image.transforms package": [[44, "module-cyto_dl.image.transforms"]], "cyto_dl.image.transforms.add_meta module": [[45, "module-cyto_dl.image.transforms.add_meta"]], "cyto_dl.image.transforms.bright_sampler module": [[46, "module-cyto_dl.image.transforms.bright_sampler"]], "cyto_dl.image.transforms.clip module": [[47, "module-cyto_dl.image.transforms.clip"]], "cyto_dl.image.transforms.contrastadjust module": [[48, "module-cyto_dl.image.transforms.contrastadjust"]], "cyto_dl.image.transforms.generate_jepa_masks module": [[49, "module-cyto_dl.image.transforms.generate_jepa_masks"]], "cyto_dl.image.transforms.merge module": [[50, "module-cyto_dl.image.transforms.merge"]], "cyto_dl.image.transforms.multiscale_cropper module": [[51, "module-cyto_dl.image.transforms.multiscale_cropper"]], "cyto_dl.image.transforms.pad module": [[52, "module-cyto_dl.image.transforms.pad"]], "cyto_dl.image.transforms.project module": [[53, "module-cyto_dl.image.transforms.project"]], "cyto_dl.image.transforms.rotation_mask_transform module": [[54, "cyto-dl-image-transforms-rotation-mask-transform-module"]], "cyto_dl.image.transforms.save module": [[55, "module-cyto_dl.image.transforms.save"]], "cyto_dl.image.transforms.so2_random_rotation module": [[56, "cyto-dl-image-transforms-so2-random-rotation-module"]], "cyto_dl.image.transforms.track_transforms module": [[57, "module-cyto_dl.image.transforms.track_transforms"]], "cyto_dl.loggers package": [[58, "module-cyto_dl.loggers"]], "cyto_dl.loggers.mlflow module": [[59, "module-cyto_dl.loggers.mlflow"]], "cyto_dl.models package": [[60, "module-cyto_dl.models"]], "cyto_dl.models.base_model module": [[61, "module-cyto_dl.models.base_model"]], "cyto_dl.models.basic_model module": [[62, "module-cyto_dl.models.basic_model"]], "cyto_dl.models.classification package": [[63, "module-cyto_dl.models.classification"]], "cyto_dl.models.classification.classification module": [[64, "module-cyto_dl.models.classification.classification"]], "cyto_dl.models.classification.timepoint_classification module": [[65, "module-cyto_dl.models.classification.timepoint_classification"]], "cyto_dl.models.contrastive package": [[66, "module-cyto_dl.models.contrastive"]], "cyto_dl.models.contrastive.contrastive module": [[67, "module-cyto_dl.models.contrastive.contrastive"]], "cyto_dl.models.handlers package": [[68, "cyto-dl-models-handlers-package"]], "cyto_dl.models.handlers.base_handler module": [[69, "cyto-dl-models-handlers-base-handler-module"]], "cyto_dl.models.handlers.image_handler module": [[70, "cyto-dl-models-handlers-image-handler-module"]], "cyto_dl.models.handlers.load_image_patch module": [[71, "cyto-dl-models-handlers-load-image-patch-module"]], "cyto_dl.models.im2im package": [[72, "module-cyto_dl.models.im2im"]], "cyto_dl.models.im2im.gan module": [[73, "module-cyto_dl.models.im2im.gan"]], "cyto_dl.models.im2im.multi_task module": [[74, "module-cyto_dl.models.im2im.multi_task"]], "cyto_dl.models.im2im.utils package": [[75, "module-cyto_dl.models.im2im.utils"]], "cyto_dl.models.im2im.utils.instance_seg module": [[76, "module-cyto_dl.models.im2im.utils.instance_seg"]], "cyto_dl.models.im2im.utils.noise_annealer module": [[77, "module-cyto_dl.models.im2im.utils.noise_annealer"]], "cyto_dl.models.im2im.utils.postprocessing package": [[78, "module-cyto_dl.models.im2im.utils.postprocessing"]], "cyto_dl.models.im2im.utils.postprocessing.act_thresh_label module": [[79, "module-cyto_dl.models.im2im.utils.postprocessing.act_thresh_label"]], "cyto_dl.models.im2im.utils.postprocessing.auto_thresh module": [[80, "module-cyto_dl.models.im2im.utils.postprocessing.auto_thresh"]], "cyto_dl.models.im2im.utils.postprocessing.dict_to_im module": [[81, "module-cyto_dl.models.im2im.utils.postprocessing.dict_to_im"]], "cyto_dl.models.jepa package": [[82, "module-cyto_dl.models.jepa"]], "cyto_dl.models.jepa.ijepa module": [[83, "module-cyto_dl.models.jepa.ijepa"]], "cyto_dl.models.jepa.iwm module": [[84, "module-cyto_dl.models.jepa.iwm"]], "cyto_dl.models.jepa.jepa_base module": [[85, "module-cyto_dl.models.jepa.jepa_base"]], "cyto_dl.models.utils package": [[86, "module-cyto_dl.models.utils"]], "cyto_dl.models.utils.mlflow module": [[87, "module-cyto_dl.models.utils.mlflow"]], "cyto_dl.models.utils.utils module": [[88, "module-cyto_dl.models.utils.utils"]], "cyto_dl.models.vae package": [[89, "cyto-dl-models-vae-package"]], "cyto_dl.models.vae.base_vae module": [[90, "module-cyto_dl.models.vae.base_vae"]], "cyto_dl.models.vae.image_canon_vae module": [[91, "cyto-dl-models-vae-image-canon-vae-module"]], "cyto_dl.models.vae.image_encoder module": [[92, "cyto-dl-models-vae-image-encoder-module"]], "cyto_dl.models.vae.image_vae module": [[93, "cyto-dl-models-vae-image-vae-module"]], "cyto_dl.models.vae.latent_loss_vae module": [[94, "cyto-dl-models-vae-latent-loss-vae-module"]], "cyto_dl.models.vae.o2_spharm_vae package": [[95, "cyto-dl-models-vae-o2-spharm-vae-package"]], "cyto_dl.models.vae.o2_spharm_vae.o2_spharm_encoder module": [[96, "cyto-dl-models-vae-o2-spharm-vae-o2-spharm-encoder-module"]], "cyto_dl.models.vae.o2_spharm_vae.o2_spharm_vae module": [[97, "cyto-dl-models-vae-o2-spharm-vae-o2-spharm-vae-module"]], "cyto_dl.models.vae.point_cloud_vae module": [[98, "cyto-dl-models-vae-point-cloud-vae-module"]], "cyto_dl.models.vae.priors package": [[99, "module-cyto_dl.models.vae.priors"]], "cyto_dl.models.vae.priors.abstract_prior module": [[100, "module-cyto_dl.models.vae.priors.abstract_prior"]], "cyto_dl.models.vae.priors.gaussian module": [[101, "module-cyto_dl.models.vae.priors.gaussian"]], "cyto_dl.models.vae.priors.identity_prior module": [[102, "module-cyto_dl.models.vae.priors.identity_prior"]], "cyto_dl.models.vae.priors.joint_prior module": [[103, "module-cyto_dl.models.vae.priors.joint_prior"]], "cyto_dl.models.vae.tabular_vae module": [[104, "cyto-dl-models-vae-tabular-vae-module"]], "cyto_dl.nn package": [[105, "module-cyto_dl.nn"]], "cyto_dl.nn.discriminators package": [[106, "module-cyto_dl.nn.discriminators"]], "cyto_dl.nn.discriminators.multi_scale_discriminator module": [[107, "module-cyto_dl.nn.discriminators.multi_scale_discriminator"]], "cyto_dl.nn.discriminators.n_layer_discriminator module": [[108, "module-cyto_dl.nn.discriminators.n_layer_discriminator"]], "cyto_dl.nn.head package": [[109, "module-cyto_dl.nn.head"]], "cyto_dl.nn.head.base_head module": [[110, "module-cyto_dl.nn.head.base_head"]], "cyto_dl.nn.head.gan_head module": [[111, "module-cyto_dl.nn.head.gan_head"]], "cyto_dl.nn.head.gan_head_superres module": [[112, "module-cyto_dl.nn.head.gan_head_superres"]], "cyto_dl.nn.head.mae_head module": [[113, "module-cyto_dl.nn.head.mae_head"]], "cyto_dl.nn.head.mask_head module": [[114, "module-cyto_dl.nn.head.mask_head"]], "cyto_dl.nn.head.res_blocks_head module": [[115, "module-cyto_dl.nn.head.res_blocks_head"]], "cyto_dl.nn.head.vic_reg module": [[116, "module-cyto_dl.nn.head.vic_reg"]], "cyto_dl.nn.hr_skip module": [[117, "module-cyto_dl.nn.hr_skip"]], "cyto_dl.nn.losses package": [[118, "module-cyto_dl.nn.losses"]], "cyto_dl.nn.losses.adversarial_loss module": [[119, "module-cyto_dl.nn.losses.adversarial_loss"]], "cyto_dl.nn.losses.chamfer_loss module": [[120, "module-cyto_dl.nn.losses.chamfer_loss"]], "cyto_dl.nn.losses.continuous_bernoulli module": [[121, "module-cyto_dl.nn.losses.continuous_bernoulli"]], "cyto_dl.nn.losses.cosine_loss module": [[122, "module-cyto_dl.nn.losses.cosine_loss"]], "cyto_dl.nn.losses.gan_loss module": [[123, "module-cyto_dl.nn.losses.gan_loss"]], "cyto_dl.nn.losses.gaussian_nll_loss module": [[124, "module-cyto_dl.nn.losses.gaussian_nll_loss"]], "cyto_dl.nn.losses.geomloss module": [[125, "cyto-dl-nn-losses-geomloss-module"]], "cyto_dl.nn.losses.loss_wrapper module": [[126, "module-cyto_dl.nn.losses.loss_wrapper"]], "cyto_dl.nn.losses.spharm_loss module": [[127, "cyto-dl-nn-losses-spharm-loss-module"]], "cyto_dl.nn.losses.threshold_loss module": [[128, "module-cyto_dl.nn.losses.threshold_loss"]], "cyto_dl.nn.losses.vic_reg module": [[129, "module-cyto_dl.nn.losses.vic_reg"]], "cyto_dl.nn.losses.weibull module": [[130, "module-cyto_dl.nn.losses.weibull"]], "cyto_dl.nn.losses.weighted_mse_loss module": [[131, "module-cyto_dl.nn.losses.weighted_mse_loss"]], "cyto_dl.nn.mlp module": [[132, "module-cyto_dl.nn.mlp"]], "cyto_dl.nn.point_cloud package": [[133, "cyto-dl-nn-point-cloud-package"]], "cyto_dl.nn.point_cloud.dgcnn module": [[134, "cyto-dl-nn-point-cloud-dgcnn-module"]], "cyto_dl.nn.point_cloud.folding_net module": [[135, "cyto-dl-nn-point-cloud-folding-net-module"]], "cyto_dl.nn.point_cloud.graph_functions module": [[136, "cyto-dl-nn-point-cloud-graph-functions-module"]], "cyto_dl.nn.point_cloud.vnn module": [[137, "cyto-dl-nn-point-cloud-vnn-module"]], "cyto_dl.nn.res_unit module": [[138, "module-cyto_dl.nn.res_unit"]], "cyto_dl.nn.spatial_transformer module": [[139, "module-cyto_dl.nn.spatial_transformer"]], "cyto_dl.nn.track_sequence_predictor module": [[140, "module-cyto_dl.nn.track_sequence_predictor"]], "cyto_dl.nn.vits package": [[141, "module-cyto_dl.nn.vits"]], "cyto_dl.nn.vits.blocks package": [[142, "module-cyto_dl.nn.vits.blocks"]], "cyto_dl.nn.vits.blocks.cross_attention module": [[143, "module-cyto_dl.nn.vits.blocks.cross_attention"]], "cyto_dl.nn.vits.blocks.intermediate_weigher module": [[144, "module-cyto_dl.nn.vits.blocks.intermediate_weigher"]], "cyto_dl.nn.vits.blocks.patchify module": [[145, "module-cyto_dl.nn.vits.blocks.patchify"]], "cyto_dl.nn.vits.cross_mae module": [[146, "module-cyto_dl.nn.vits.cross_mae"]], "cyto_dl.nn.vits.jepa module": [[147, "module-cyto_dl.nn.vits.jepa"]], "cyto_dl.nn.vits.mae module": [[148, "module-cyto_dl.nn.vits.mae"]], "cyto_dl.nn.vits.seg module": [[149, "module-cyto_dl.nn.vits.seg"]], "cyto_dl.nn.vits.utils module": [[150, "module-cyto_dl.nn.vits.utils"]], "cyto_dl.point_cloud package": [[151, "module-cyto_dl.point_cloud"]], "cyto_dl.point_cloud.io package": [[152, "cyto-dl-point-cloud-io-package"]], "cyto_dl.point_cloud.io.read_pcloud module": [[153, "cyto-dl-point-cloud-io-read-pcloud-module"]], "cyto_dl.train module": [[154, "module-cyto_dl.train"]], "cyto_dl.utils package": [[155, "module-cyto_dl.utils"]], "cyto_dl.utils.array module": [[156, "module-cyto_dl.utils.array"]], "cyto_dl.utils.checkpoint module": [[157, "module-cyto_dl.utils.checkpoint"]], "cyto_dl.utils.config module": [[158, "module-cyto_dl.utils.config"]], "cyto_dl.utils.download_test_data module": [[159, "module-cyto_dl.utils.download_test_data"]], "cyto_dl.utils.dummy_dataset module": [[160, "module-cyto_dl.utils.dummy_dataset"]], "cyto_dl.utils.pylogger module": [[161, "module-cyto_dl.utils.pylogger"]], "cyto_dl.utils.rich_utils module": [[162, "module-cyto_dl.utils.rich_utils"]], "cyto_dl.utils.rotation module": [[163, "module-cyto_dl.utils.rotation"]], "cyto_dl.utils.spharm package": [[164, "cyto-dl-utils-spharm-package"]], "cyto_dl.utils.spharm.reconstruction module": [[165, "cyto-dl-utils-spharm-reconstruction-module"]], "cyto_dl.utils.spharm.rotation module": [[166, "cyto-dl-utils-spharm-rotation-module"]], "cyto_dl.utils.template_utils module": [[167, "module-cyto_dl.utils.template_utils"]], "Overview": [[168, "overview"], [171, "overview"]], "data": [[168, "data"], [171, "data"]], "experiment": [[168, "experiment"], [171, "experiment"]], "model": [[168, "model"], [171, "model"]], "Installation": [[169, "installation"]], "cyto_dl": [[170, "cyto-dl"]], "Modifying the Example Configs": [[172, "modifying-the-example-configs"]], "Training": [[172, "training"]], "Testing/Prediction": [[172, "testing-prediction"]]}, "indexentries": {"cyto_dl": [[1, "module-cyto_dl"]], "module": [[1, "module-cyto_dl"], [2, "module-cyto_dl.api"], [3, "module-cyto_dl.api.cyto_dl_model"], [4, "module-cyto_dl.api.cyto_dl_model.cyto_dl_base_model"], [5, "module-cyto_dl.api.cyto_dl_model.segmentation_plugin_model"], [6, "module-cyto_dl.api.data"], [7, "module-cyto_dl.api.model"], [8, "module-cyto_dl.callbacks"], [10, "module-cyto_dl.callbacks.csv_saver"], [11, "module-cyto_dl.callbacks.image_saver"], [13, "module-cyto_dl.callbacks.layer_freeze"], [14, "module-cyto_dl.callbacks.model_utils"], [15, "module-cyto_dl.callbacks.outlier_detection"], [17, "module-cyto_dl.dataframe"], [18, "module-cyto_dl.dataframe.readers"], [19, "module-cyto_dl.dataframe.transforms"], [20, "module-cyto_dl.dataframe.transforms.filter"], [21, "module-cyto_dl.dataframe.transforms.group_cols"], [22, "module-cyto_dl.dataframe.transforms.misc"], [23, "module-cyto_dl.dataframe.transforms.split"], [24, "module-cyto_dl.datamodules"], [25, "module-cyto_dl.datamodules.array"], [26, "module-cyto_dl.datamodules.data_dict"], [27, "module-cyto_dl.datamodules.dataframe"], [28, "module-cyto_dl.datamodules.dataframe.dataframe_datamodule"], [29, "module-cyto_dl.datamodules.dataframe.grouped_dataframe_datamodule"], [30, "module-cyto_dl.datamodules.dataframe.utils"], [31, "module-cyto_dl.datamodules.folder"], [32, "module-cyto_dl.datamodules.multidim_image"], [33, "module-cyto_dl.datamodules.smartcache"], [34, "module-cyto_dl.datamodules.torchvision"], [35, "module-cyto_dl.eval"], [36, "module-cyto_dl.image"], [37, "module-cyto_dl.image.io"], [38, "module-cyto_dl.image.io.aicsimage_loader"], [39, "module-cyto_dl.image.io.monai_bio_reader"], [40, "module-cyto_dl.image.io.numpy_reader"], [41, "module-cyto_dl.image.io.ome_zarr_reader"], [42, "module-cyto_dl.image.io.polygon_loader"], [43, "module-cyto_dl.image.io.skimage_reader"], [44, "module-cyto_dl.image.transforms"], [45, "module-cyto_dl.image.transforms.add_meta"], [46, "module-cyto_dl.image.transforms.bright_sampler"], [47, "module-cyto_dl.image.transforms.clip"], [48, "module-cyto_dl.image.transforms.contrastadjust"], [49, "module-cyto_dl.image.transforms.generate_jepa_masks"], [50, "module-cyto_dl.image.transforms.merge"], [51, "module-cyto_dl.image.transforms.multiscale_cropper"], [52, "module-cyto_dl.image.transforms.pad"], [53, "module-cyto_dl.image.transforms.project"], [55, "module-cyto_dl.image.transforms.save"], [57, "module-cyto_dl.image.transforms.track_transforms"], [58, "module-cyto_dl.loggers"], [59, "module-cyto_dl.loggers.mlflow"], [60, "module-cyto_dl.models"], [61, "module-cyto_dl.models.base_model"], [62, "module-cyto_dl.models.basic_model"], [63, "module-cyto_dl.models.classification"], [64, "module-cyto_dl.models.classification.classification"], [65, "module-cyto_dl.models.classification.timepoint_classification"], [66, "module-cyto_dl.models.contrastive"], [67, "module-cyto_dl.models.contrastive.contrastive"], [72, "module-cyto_dl.models.im2im"], [73, "module-cyto_dl.models.im2im.gan"], [74, "module-cyto_dl.models.im2im.multi_task"], [75, "module-cyto_dl.models.im2im.utils"], [76, "module-cyto_dl.models.im2im.utils.instance_seg"], [77, "module-cyto_dl.models.im2im.utils.noise_annealer"], [78, "module-cyto_dl.models.im2im.utils.postprocessing"], [79, "module-cyto_dl.models.im2im.utils.postprocessing.act_thresh_label"], [80, "module-cyto_dl.models.im2im.utils.postprocessing.auto_thresh"], [81, "module-cyto_dl.models.im2im.utils.postprocessing.dict_to_im"], [82, "module-cyto_dl.models.jepa"], [83, "module-cyto_dl.models.jepa.ijepa"], [84, "module-cyto_dl.models.jepa.iwm"], [85, "module-cyto_dl.models.jepa.jepa_base"], [86, "module-cyto_dl.models.utils"], [87, "module-cyto_dl.models.utils.mlflow"], [88, "module-cyto_dl.models.utils.utils"], [90, "module-cyto_dl.models.vae.base_vae"], [99, "module-cyto_dl.models.vae.priors"], [100, "module-cyto_dl.models.vae.priors.abstract_prior"], [101, "module-cyto_dl.models.vae.priors.gaussian"], [102, "module-cyto_dl.models.vae.priors.identity_prior"], [103, "module-cyto_dl.models.vae.priors.joint_prior"], [105, "module-cyto_dl.nn"], [106, "module-cyto_dl.nn.discriminators"], [107, "module-cyto_dl.nn.discriminators.multi_scale_discriminator"], [108, "module-cyto_dl.nn.discriminators.n_layer_discriminator"], [109, "module-cyto_dl.nn.head"], [110, "module-cyto_dl.nn.head.base_head"], [111, "module-cyto_dl.nn.head.gan_head"], [112, "module-cyto_dl.nn.head.gan_head_superres"], [113, "module-cyto_dl.nn.head.mae_head"], [114, "module-cyto_dl.nn.head.mask_head"], [115, "module-cyto_dl.nn.head.res_blocks_head"], [116, "module-cyto_dl.nn.head.vic_reg"], [117, "module-cyto_dl.nn.hr_skip"], [118, "module-cyto_dl.nn.losses"], [119, "module-cyto_dl.nn.losses.adversarial_loss"], [120, "module-cyto_dl.nn.losses.chamfer_loss"], [121, "module-cyto_dl.nn.losses.continuous_bernoulli"], [122, "module-cyto_dl.nn.losses.cosine_loss"], [123, "module-cyto_dl.nn.losses.gan_loss"], [124, "module-cyto_dl.nn.losses.gaussian_nll_loss"], [126, "module-cyto_dl.nn.losses.loss_wrapper"], [128, "module-cyto_dl.nn.losses.threshold_loss"], [129, "module-cyto_dl.nn.losses.vic_reg"], [130, "module-cyto_dl.nn.losses.weibull"], [131, "module-cyto_dl.nn.losses.weighted_mse_loss"], [132, "module-cyto_dl.nn.mlp"], [138, "module-cyto_dl.nn.res_unit"], [139, "module-cyto_dl.nn.spatial_transformer"], [140, "module-cyto_dl.nn.track_sequence_predictor"], [141, "module-cyto_dl.nn.vits"], [142, "module-cyto_dl.nn.vits.blocks"], [143, "module-cyto_dl.nn.vits.blocks.cross_attention"], [144, "module-cyto_dl.nn.vits.blocks.intermediate_weigher"], [145, "module-cyto_dl.nn.vits.blocks.patchify"], [146, "module-cyto_dl.nn.vits.cross_mae"], [147, "module-cyto_dl.nn.vits.jepa"], [148, "module-cyto_dl.nn.vits.mae"], [149, "module-cyto_dl.nn.vits.seg"], [150, "module-cyto_dl.nn.vits.utils"], [151, "module-cyto_dl.point_cloud"], [154, "module-cyto_dl.train"], [155, "module-cyto_dl.utils"], [156, "module-cyto_dl.utils.array"], [157, "module-cyto_dl.utils.checkpoint"], [158, "module-cyto_dl.utils.config"], [159, "module-cyto_dl.utils.download_test_data"], [160, "module-cyto_dl.utils.dummy_dataset"], [161, "module-cyto_dl.utils.pylogger"], [162, "module-cyto_dl.utils.rich_utils"], [163, "module-cyto_dl.utils.rotation"], [167, "module-cyto_dl.utils.template_utils"]], "cyto_dl.api": [[2, "module-cyto_dl.api"]], "cyto_dl.api.cyto_dl_model": [[3, "module-cyto_dl.api.cyto_dl_model"]], "cytodlbasemodel (class in cyto_dl.api.cyto_dl_model.cyto_dl_base_model)": [[4, "cyto_dl.api.cyto_dl_model.cyto_dl_base_model.CytoDLBaseModel"]], "cyto_dl.api.cyto_dl_model.cyto_dl_base_model": [[4, "module-cyto_dl.api.cyto_dl_model.cyto_dl_base_model"]], "from_default_config() (cyto_dl.api.cyto_dl_model.cyto_dl_base_model.cytodlbasemodel class method)": [[4, "cyto_dl.api.cyto_dl_model.cyto_dl_base_model.CytoDLBaseModel.from_default_config"]], "from_existing_config() (cyto_dl.api.cyto_dl_model.cyto_dl_base_model.cytodlbasemodel class method)": [[4, "cyto_dl.api.cyto_dl_model.cyto_dl_base_model.CytoDLBaseModel.from_existing_config"]], "get_config() (cyto_dl.api.cyto_dl_model.cyto_dl_base_model.cytodlbasemodel method)": [[4, "cyto_dl.api.cyto_dl_model.cyto_dl_base_model.CytoDLBaseModel.get_config"]], "get_experiment_name() (cyto_dl.api.cyto_dl_model.cyto_dl_base_model.cytodlbasemodel method)": [[4, "cyto_dl.api.cyto_dl_model.cyto_dl_base_model.CytoDLBaseModel.get_experiment_name"]], "predict() (cyto_dl.api.cyto_dl_model.cyto_dl_base_model.cytodlbasemodel method)": [[4, "cyto_dl.api.cyto_dl_model.cyto_dl_base_model.CytoDLBaseModel.predict"]], "save_config() (cyto_dl.api.cyto_dl_model.cyto_dl_base_model.cytodlbasemodel method)": [[4, "cyto_dl.api.cyto_dl_model.cyto_dl_base_model.CytoDLBaseModel.save_config"]], "set_experiment_name() (cyto_dl.api.cyto_dl_model.cyto_dl_base_model.cytodlbasemodel method)": [[4, "cyto_dl.api.cyto_dl_model.cyto_dl_base_model.CytoDLBaseModel.set_experiment_name"]], "train() (cyto_dl.api.cyto_dl_model.cyto_dl_base_model.cytodlbasemodel method)": [[4, "cyto_dl.api.cyto_dl_model.cyto_dl_base_model.CytoDLBaseModel.train"]], "segmentationpluginmodel (class in cyto_dl.api.cyto_dl_model.segmentation_plugin_model)": [[5, "cyto_dl.api.cyto_dl_model.segmentation_plugin_model.SegmentationPluginModel"]], "cyto_dl.api.cyto_dl_model.segmentation_plugin_model": [[5, "module-cyto_dl.api.cyto_dl_model.segmentation_plugin_model"]], "get_hardware_type() (cyto_dl.api.cyto_dl_model.segmentation_plugin_model.segmentationpluginmodel method)": [[5, "cyto_dl.api.cyto_dl_model.segmentation_plugin_model.SegmentationPluginModel.get_hardware_type"]], "get_input_channel() (cyto_dl.api.cyto_dl_model.segmentation_plugin_model.segmentationpluginmodel method)": [[5, "cyto_dl.api.cyto_dl_model.segmentation_plugin_model.SegmentationPluginModel.get_input_channel"]], "get_manifest_column_names() (cyto_dl.api.cyto_dl_model.segmentation_plugin_model.segmentationpluginmodel method)": [[5, "cyto_dl.api.cyto_dl_model.segmentation_plugin_model.SegmentationPluginModel.get_manifest_column_names"]], "get_patch_size() (cyto_dl.api.cyto_dl_model.segmentation_plugin_model.segmentationpluginmodel method)": [[5, "cyto_dl.api.cyto_dl_model.segmentation_plugin_model.SegmentationPluginModel.get_patch_size"]], "get_raw_image_channels() (cyto_dl.api.cyto_dl_model.segmentation_plugin_model.segmentationpluginmodel method)": [[5, "cyto_dl.api.cyto_dl_model.segmentation_plugin_model.SegmentationPluginModel.get_raw_image_channels"]], "get_split_column() (cyto_dl.api.cyto_dl_model.segmentation_plugin_model.segmentationpluginmodel method)": [[5, "cyto_dl.api.cyto_dl_model.segmentation_plugin_model.SegmentationPluginModel.get_split_column"]], "remove_split_column() (cyto_dl.api.cyto_dl_model.segmentation_plugin_model.segmentationpluginmodel method)": [[5, "cyto_dl.api.cyto_dl_model.segmentation_plugin_model.SegmentationPluginModel.remove_split_column"]], "set_hardware_type() (cyto_dl.api.cyto_dl_model.segmentation_plugin_model.segmentationpluginmodel method)": [[5, "cyto_dl.api.cyto_dl_model.segmentation_plugin_model.SegmentationPluginModel.set_hardware_type"]], "set_input_channel() (cyto_dl.api.cyto_dl_model.segmentation_plugin_model.segmentationpluginmodel method)": [[5, "cyto_dl.api.cyto_dl_model.segmentation_plugin_model.SegmentationPluginModel.set_input_channel"]], "set_manifest_column_names() (cyto_dl.api.cyto_dl_model.segmentation_plugin_model.segmentationpluginmodel method)": [[5, "cyto_dl.api.cyto_dl_model.segmentation_plugin_model.SegmentationPluginModel.set_manifest_column_names"]], "set_patch_size() (cyto_dl.api.cyto_dl_model.segmentation_plugin_model.segmentationpluginmodel method)": [[5, "cyto_dl.api.cyto_dl_model.segmentation_plugin_model.SegmentationPluginModel.set_patch_size"]], "set_raw_image_channels() (cyto_dl.api.cyto_dl_model.segmentation_plugin_model.segmentationpluginmodel method)": [[5, "cyto_dl.api.cyto_dl_model.segmentation_plugin_model.SegmentationPluginModel.set_raw_image_channels"]], "set_split_column() (cyto_dl.api.cyto_dl_model.segmentation_plugin_model.segmentationpluginmodel method)": [[5, "cyto_dl.api.cyto_dl_model.segmentation_plugin_model.SegmentationPluginModel.set_split_column"]], "autothresholdmethod (class in cyto_dl.api.data)": [[6, "cyto_dl.api.data.AutoThresholdMethod"]], "cpu (cyto_dl.api.data.hardwaretype attribute)": [[6, "cyto_dl.api.data.HardwareType.CPU"]], "experimenttype (class in cyto_dl.api.data)": [[6, "cyto_dl.api.data.ExperimentType"]], "gan (cyto_dl.api.data.experimenttype attribute)": [[6, "cyto_dl.api.data.ExperimentType.GAN"]], "gpu (cyto_dl.api.data.hardwaretype attribute)": [[6, "cyto_dl.api.data.HardwareType.GPU"]], "hardwaretype (class in cyto_dl.api.data)": [[6, "cyto_dl.api.data.HardwareType"]], "instance_seg (cyto_dl.api.data.experimenttype attribute)": [[6, "cyto_dl.api.data.ExperimentType.INSTANCE_SEG"]], "isodata (cyto_dl.api.data.autothresholdmethod attribute)": [[6, "cyto_dl.api.data.AutoThresholdMethod.ISODATA"]], "label_free (cyto_dl.api.data.experimenttype attribute)": [[6, "cyto_dl.api.data.ExperimentType.LABEL_FREE"]], "large (cyto_dl.api.data.patchsize attribute)": [[6, "cyto_dl.api.data.PatchSize.LARGE"]], "li (cyto_dl.api.data.autothresholdmethod attribute)": [[6, "cyto_dl.api.data.AutoThresholdMethod.LI"]], "local (cyto_dl.api.data.autothresholdmethod attribute)": [[6, "cyto_dl.api.data.AutoThresholdMethod.LOCAL"]], "mean (cyto_dl.api.data.autothresholdmethod attribute)": [[6, "cyto_dl.api.data.AutoThresholdMethod.MEAN"]], "medium (cyto_dl.api.data.patchsize attribute)": [[6, "cyto_dl.api.data.PatchSize.MEDIUM"]], "minimum (cyto_dl.api.data.autothresholdmethod attribute)": [[6, "cyto_dl.api.data.AutoThresholdMethod.MINIMUM"]], "multiotsu (cyto_dl.api.data.autothresholdmethod attribute)": [[6, "cyto_dl.api.data.AutoThresholdMethod.MULTIOTSU"]], "niblack (cyto_dl.api.data.autothresholdmethod attribute)": [[6, "cyto_dl.api.data.AutoThresholdMethod.NIBLACK"]], "otsu (cyto_dl.api.data.autothresholdmethod attribute)": [[6, "cyto_dl.api.data.AutoThresholdMethod.OTSU"]], "patchsize (class in cyto_dl.api.data)": [[6, "cyto_dl.api.data.PatchSize"]], "sauvola (cyto_dl.api.data.autothresholdmethod attribute)": [[6, "cyto_dl.api.data.AutoThresholdMethod.SAUVOLA"]], "segmentation (cyto_dl.api.data.experimenttype attribute)": [[6, "cyto_dl.api.data.ExperimentType.SEGMENTATION"]], "segmentation_array (cyto_dl.api.data.experimenttype attribute)": [[6, "cyto_dl.api.data.ExperimentType.SEGMENTATION_ARRAY"]], "segmentation_plugin (cyto_dl.api.data.experimenttype attribute)": [[6, "cyto_dl.api.data.ExperimentType.SEGMENTATION_PLUGIN"]], "small (cyto_dl.api.data.patchsize attribute)": [[6, "cyto_dl.api.data.PatchSize.SMALL"]], "triangle (cyto_dl.api.data.autothresholdmethod attribute)": [[6, "cyto_dl.api.data.AutoThresholdMethod.TRIANGLE"]], "yen (cyto_dl.api.data.autothresholdmethod attribute)": [[6, "cyto_dl.api.data.AutoThresholdMethod.YEN"]], "cyto_dl.api.data": [[6, "module-cyto_dl.api.data"]], "cytodlmodel (class in cyto_dl.api.model)": [[7, "cyto_dl.api.model.CytoDLModel"]], "cyto_dl.api.model": [[7, "module-cyto_dl.api.model"]], "download_example_data() (cyto_dl.api.model.cytodlmodel method)": [[7, "cyto_dl.api.model.CytoDLModel.download_example_data"]], "load_config_from_dict() (cyto_dl.api.model.cytodlmodel method)": [[7, "cyto_dl.api.model.CytoDLModel.load_config_from_dict"]], "load_config_from_file() (cyto_dl.api.model.cytodlmodel method)": [[7, "cyto_dl.api.model.CytoDLModel.load_config_from_file"]], "load_default_experiment() (cyto_dl.api.model.cytodlmodel method)": [[7, "cyto_dl.api.model.CytoDLModel.load_default_experiment"]], "override_config() (cyto_dl.api.model.cytodlmodel method)": [[7, "cyto_dl.api.model.CytoDLModel.override_config"]], "predict() (cyto_dl.api.model.cytodlmodel method)": [[7, "cyto_dl.api.model.CytoDLModel.predict"]], "print_config() (cyto_dl.api.model.cytodlmodel method)": [[7, "cyto_dl.api.model.CytoDLModel.print_config"]], "save_config() (cyto_dl.api.model.cytodlmodel method)": [[7, "cyto_dl.api.model.CytoDLModel.save_config"]], "train() (cyto_dl.api.model.cytodlmodel method)": [[7, "cyto_dl.api.model.CytoDLModel.train"]], "cyto_dl.callbacks": [[8, "module-cyto_dl.callbacks"]], "csvsaver (class in cyto_dl.callbacks.csv_saver)": [[10, "cyto_dl.callbacks.csv_saver.CSVSaver"]], "jepasaver (class in cyto_dl.callbacks.csv_saver)": [[10, "cyto_dl.callbacks.csv_saver.JEPASaver"]], "cyto_dl.callbacks.csv_saver": [[10, "module-cyto_dl.callbacks.csv_saver"]], "on_predict_epoch_end() (cyto_dl.callbacks.csv_saver.csvsaver method)": [[10, "cyto_dl.callbacks.csv_saver.CSVSaver.on_predict_epoch_end"]], "pred_to_csv() (cyto_dl.callbacks.csv_saver.csvsaver method)": [[10, "cyto_dl.callbacks.csv_saver.CSVSaver.pred_to_csv"]], "pred_to_csv() (cyto_dl.callbacks.csv_saver.jepasaver method)": [[10, "cyto_dl.callbacks.csv_saver.JEPASaver.pred_to_csv"]], "save_feats() (cyto_dl.callbacks.csv_saver.csvsaver method)": [[10, "cyto_dl.callbacks.csv_saver.CSVSaver.save_feats"]], "imagesaver (class in cyto_dl.callbacks.image_saver)": [[11, "cyto_dl.callbacks.image_saver.ImageSaver"]], "cyto_dl.callbacks.image_saver": [[11, "module-cyto_dl.callbacks.image_saver"]], "on_predict_batch_end() (cyto_dl.callbacks.image_saver.imagesaver method)": [[11, "cyto_dl.callbacks.image_saver.ImageSaver.on_predict_batch_end"]], "on_test_batch_end() (cyto_dl.callbacks.image_saver.imagesaver method)": [[11, "cyto_dl.callbacks.image_saver.ImageSaver.on_test_batch_end"]], "on_train_batch_end() (cyto_dl.callbacks.image_saver.imagesaver method)": [[11, "cyto_dl.callbacks.image_saver.ImageSaver.on_train_batch_end"]], "on_validation_batch_end() (cyto_dl.callbacks.image_saver.imagesaver method)": [[11, "cyto_dl.callbacks.image_saver.ImageSaver.on_validation_batch_end"]], "save() (cyto_dl.callbacks.image_saver.imagesaver method)": [[11, "cyto_dl.callbacks.image_saver.ImageSaver.save"]], "layerfreeze (class in cyto_dl.callbacks.layer_freeze)": [[13, "cyto_dl.callbacks.layer_freeze.LayerFreeze"]], "cyto_dl.callbacks.layer_freeze": [[13, "module-cyto_dl.callbacks.layer_freeze"]], "setup() (cyto_dl.callbacks.layer_freeze.layerfreeze method)": [[13, "cyto_dl.callbacks.layer_freeze.LayerFreeze.setup"]], "getembeddings (class in cyto_dl.callbacks.model_utils)": [[14, "cyto_dl.callbacks.model_utils.GetEmbeddings"]], "cyto_dl.callbacks.model_utils": [[14, "module-cyto_dl.callbacks.model_utils"]], "get_all_embeddings() (in module cyto_dl.callbacks.model_utils)": [[14, "cyto_dl.callbacks.model_utils.get_all_embeddings"]], "on_test_epoch_end() (cyto_dl.callbacks.model_utils.getembeddings method)": [[14, "cyto_dl.callbacks.model_utils.GetEmbeddings.on_test_epoch_end"]], "save_predictions_classifier() (in module cyto_dl.callbacks.model_utils)": [[14, "cyto_dl.callbacks.model_utils.save_predictions_classifier"]], "outlierdetection (class in cyto_dl.callbacks.outlier_detection)": [[15, "cyto_dl.callbacks.outlier_detection.OutlierDetection"]], "calculate_mahalanobis_hook() (cyto_dl.callbacks.outlier_detection.outlierdetection method)": [[15, "cyto_dl.callbacks.outlier_detection.OutlierDetection.calculate_mahalanobis_hook"]], "cyto_dl.callbacks.outlier_detection": [[15, "module-cyto_dl.callbacks.outlier_detection"]], "flatten_activations() (cyto_dl.callbacks.outlier_detection.outlierdetection method)": [[15, "cyto_dl.callbacks.outlier_detection.OutlierDetection.flatten_activations"]], "on_load_checkpoint() (cyto_dl.callbacks.outlier_detection.outlierdetection method)": [[15, "cyto_dl.callbacks.outlier_detection.OutlierDetection.on_load_checkpoint"]], "on_predict_batch_end() (cyto_dl.callbacks.outlier_detection.outlierdetection method)": [[15, "cyto_dl.callbacks.outlier_detection.OutlierDetection.on_predict_batch_end"]], "on_predict_epoch_start() (cyto_dl.callbacks.outlier_detection.outlierdetection method)": [[15, "cyto_dl.callbacks.outlier_detection.OutlierDetection.on_predict_epoch_start"]], "on_save_checkpoint() (cyto_dl.callbacks.outlier_detection.outlierdetection method)": [[15, "cyto_dl.callbacks.outlier_detection.OutlierDetection.on_save_checkpoint"]], "on_test_batch_end() (cyto_dl.callbacks.outlier_detection.outlierdetection method)": [[15, "cyto_dl.callbacks.outlier_detection.OutlierDetection.on_test_batch_end"]], "on_test_epoch_start() (cyto_dl.callbacks.outlier_detection.outlierdetection method)": [[15, "cyto_dl.callbacks.outlier_detection.OutlierDetection.on_test_epoch_start"]], "on_train_epoch_start() (cyto_dl.callbacks.outlier_detection.outlierdetection method)": [[15, "cyto_dl.callbacks.outlier_detection.OutlierDetection.on_train_epoch_start"]], "update_covariance_hook() (cyto_dl.callbacks.outlier_detection.outlierdetection method)": [[15, "cyto_dl.callbacks.outlier_detection.OutlierDetection.update_covariance_hook"]], "cyto_dl.dataframe": [[17, "module-cyto_dl.dataframe"]], "cyto_dl.dataframe.readers": [[18, "module-cyto_dl.dataframe.readers"]], "filter_columns() (in module cyto_dl.dataframe.readers)": [[18, "cyto_dl.dataframe.readers.filter_columns"]], "read_csv() (in module cyto_dl.dataframe.readers)": [[18, "cyto_dl.dataframe.readers.read_csv"]], "read_dataframe() (in module cyto_dl.dataframe.readers)": [[18, "cyto_dl.dataframe.readers.read_dataframe"]], "read_h5ad() (in module cyto_dl.dataframe.readers)": [[18, "cyto_dl.dataframe.readers.read_h5ad"]], "read_parquet() (in module cyto_dl.dataframe.readers)": [[18, "cyto_dl.dataframe.readers.read_parquet"]], "cyto_dl.dataframe.transforms": [[19, "module-cyto_dl.dataframe.transforms"]], "cyto_dl.dataframe.transforms.filter": [[20, "module-cyto_dl.dataframe.transforms.filter"]], "filter_columns() (in module cyto_dl.dataframe.transforms.filter)": [[20, "cyto_dl.dataframe.transforms.filter.filter_columns"]], "filter_rows() (in module cyto_dl.dataframe.transforms.filter)": [[20, "cyto_dl.dataframe.transforms.filter.filter_rows"]], "groupcols (class in cyto_dl.dataframe.transforms.group_cols)": [[21, "cyto_dl.dataframe.transforms.group_cols.GroupCols"]], "cyto_dl.dataframe.transforms.group_cols": [[21, "module-cyto_dl.dataframe.transforms.group_cols"]], "append_class_weights() (in module cyto_dl.dataframe.transforms.misc)": [[22, "cyto_dl.dataframe.transforms.misc.append_class_weights"]], "append_labels_to_integers() (in module cyto_dl.dataframe.transforms.misc)": [[22, "cyto_dl.dataframe.transforms.misc.append_labels_to_integers"]], "append_one_hot() (in module cyto_dl.dataframe.transforms.misc)": [[22, "cyto_dl.dataframe.transforms.misc.append_one_hot"]], "cyto_dl.dataframe.transforms.misc": [[22, "module-cyto_dl.dataframe.transforms.misc"]], "make_random_df() (in module cyto_dl.dataframe.transforms.misc)": [[22, "cyto_dl.dataframe.transforms.misc.make_random_df"]], "cyto_dl.dataframe.transforms.split": [[23, "module-cyto_dl.dataframe.transforms.split"]], "sample_n_each() (in module cyto_dl.dataframe.transforms.split)": [[23, "cyto_dl.dataframe.transforms.split.sample_n_each"]], "split_dataframe() (in module cyto_dl.dataframe.transforms.split)": [[23, "cyto_dl.dataframe.transforms.split.split_dataframe"]], "cyto_dl.datamodules": [[24, "module-cyto_dl.datamodules"]], "cyto_dl.datamodules.array": [[25, "module-cyto_dl.datamodules.array"]], "make_array_dataloader() (in module cyto_dl.datamodules.array)": [[25, "cyto_dl.datamodules.array.make_array_dataloader"]], "cyto_dl.datamodules.data_dict": [[26, "module-cyto_dl.datamodules.data_dict"]], "make_data_dict_dataloader() (in module cyto_dl.datamodules.data_dict)": [[26, "cyto_dl.datamodules.data_dict.make_data_dict_dataloader"]], "cyto_dl.datamodules.dataframe": [[27, "module-cyto_dl.datamodules.dataframe"]], "dataframedatamodule (class in cyto_dl.datamodules.dataframe.dataframe_datamodule)": [[28, "cyto_dl.datamodules.dataframe.dataframe_datamodule.DataframeDatamodule"]], "cyto_dl.datamodules.dataframe.dataframe_datamodule": [[28, "module-cyto_dl.datamodules.dataframe.dataframe_datamodule"]], "get_dataloader() (cyto_dl.datamodules.dataframe.dataframe_datamodule.dataframedatamodule method)": [[28, "cyto_dl.datamodules.dataframe.dataframe_datamodule.DataframeDatamodule.get_dataloader"]], "get_dataset() (cyto_dl.datamodules.dataframe.dataframe_datamodule.dataframedatamodule method)": [[28, "cyto_dl.datamodules.dataframe.dataframe_datamodule.DataframeDatamodule.get_dataset"]], "make_dataloader() (cyto_dl.datamodules.dataframe.dataframe_datamodule.dataframedatamodule method)": [[28, "cyto_dl.datamodules.dataframe.dataframe_datamodule.DataframeDatamodule.make_dataloader"]], "predict_dataloader() (cyto_dl.datamodules.dataframe.dataframe_datamodule.dataframedatamodule method)": [[28, "cyto_dl.datamodules.dataframe.dataframe_datamodule.DataframeDatamodule.predict_dataloader"]], "test_dataloader() (cyto_dl.datamodules.dataframe.dataframe_datamodule.dataframedatamodule method)": [[28, "cyto_dl.datamodules.dataframe.dataframe_datamodule.DataframeDatamodule.test_dataloader"]], "train_dataloader() (cyto_dl.datamodules.dataframe.dataframe_datamodule.dataframedatamodule method)": [[28, "cyto_dl.datamodules.dataframe.dataframe_datamodule.DataframeDatamodule.train_dataloader"]], "val_dataloader() (cyto_dl.datamodules.dataframe.dataframe_datamodule.dataframedatamodule method)": [[28, "cyto_dl.datamodules.dataframe.dataframe_datamodule.DataframeDatamodule.val_dataloader"]], "groupeddataframedatamodule (class in cyto_dl.datamodules.dataframe.grouped_dataframe_datamodule)": [[29, "cyto_dl.datamodules.dataframe.grouped_dataframe_datamodule.GroupedDataframeDatamodule"]], "cyto_dl.datamodules.dataframe.grouped_dataframe_datamodule": [[29, "module-cyto_dl.datamodules.dataframe.grouped_dataframe_datamodule"]], "make_dataloader() (cyto_dl.datamodules.dataframe.grouped_dataframe_datamodule.groupeddataframedatamodule method)": [[29, "cyto_dl.datamodules.dataframe.grouped_dataframe_datamodule.GroupedDataframeDatamodule.make_dataloader"]], "alternatingbatchsampler (class in cyto_dl.datamodules.dataframe.utils)": [[30, "cyto_dl.datamodules.dataframe.utils.AlternatingBatchSampler"]], "removenankeysd (class in cyto_dl.datamodules.dataframe.utils)": [[30, "cyto_dl.datamodules.dataframe.utils.RemoveNaNKeysd"]], "cyto_dl.datamodules.dataframe.utils": [[30, "module-cyto_dl.datamodules.dataframe.utils"]], "get_canonical_split_name() (in module cyto_dl.datamodules.dataframe.utils)": [[30, "cyto_dl.datamodules.dataframe.utils.get_canonical_split_name"]], "get_dataset() (in module cyto_dl.datamodules.dataframe.utils)": [[30, "cyto_dl.datamodules.dataframe.utils.get_dataset"]], "make_multiple_dataframe_splits() (in module cyto_dl.datamodules.dataframe.utils)": [[30, "cyto_dl.datamodules.dataframe.utils.make_multiple_dataframe_splits"]], "make_single_dataframe_splits() (in module cyto_dl.datamodules.dataframe.utils)": [[30, "cyto_dl.datamodules.dataframe.utils.make_single_dataframe_splits"]], "parse_transforms() (in module cyto_dl.datamodules.dataframe.utils)": [[30, "cyto_dl.datamodules.dataframe.utils.parse_transforms"]], "cyto_dl.datamodules.folder": [[31, "module-cyto_dl.datamodules.folder"]], "make_folder_dataloader() (in module cyto_dl.datamodules.folder)": [[31, "cyto_dl.datamodules.folder.make_folder_dataloader"]], "multidimimagedataset (class in cyto_dl.datamodules.multidim_image)": [[32, "cyto_dl.datamodules.multidim_image.MultiDimImageDataset"]], "create_metatensor() (cyto_dl.datamodules.multidim_image.multidimimagedataset method)": [[32, "cyto_dl.datamodules.multidim_image.MultiDimImageDataset.create_metatensor"]], "cyto_dl.datamodules.multidim_image": [[32, "module-cyto_dl.datamodules.multidim_image"]], "get_per_file_args() (cyto_dl.datamodules.multidim_image.multidimimagedataset method)": [[32, "cyto_dl.datamodules.multidim_image.MultiDimImageDataset.get_per_file_args"]], "is_batch() (cyto_dl.datamodules.multidim_image.multidimimagedataset method)": [[32, "cyto_dl.datamodules.multidim_image.MultiDimImageDataset.is_batch"]], "make_multidim_image_dataloader() (in module cyto_dl.datamodules.multidim_image)": [[32, "cyto_dl.datamodules.multidim_image.make_multidim_image_dataloader"]], "smartcachedatamodule (class in cyto_dl.datamodules.smartcache)": [[33, "cyto_dl.datamodules.smartcache.SmartcacheDatamodule"]], "cyto_dl.datamodules.smartcache": [[33, "module-cyto_dl.datamodules.smartcache"]], "get_per_file_args() (cyto_dl.datamodules.smartcache.smartcachedatamodule method)": [[33, "cyto_dl.datamodules.smartcache.SmartcacheDatamodule.get_per_file_args"]], "make_dataloader() (cyto_dl.datamodules.smartcache.smartcachedatamodule method)": [[33, "cyto_dl.datamodules.smartcache.SmartcacheDatamodule.make_dataloader"]], "predict_dataloader() (cyto_dl.datamodules.smartcache.smartcachedatamodule method)": [[33, "cyto_dl.datamodules.smartcache.SmartcacheDatamodule.predict_dataloader"]], "prepare_data() (cyto_dl.datamodules.smartcache.smartcachedatamodule method)": [[33, "cyto_dl.datamodules.smartcache.SmartcacheDatamodule.prepare_data"]], "setup() (cyto_dl.datamodules.smartcache.smartcachedatamodule method)": [[33, "cyto_dl.datamodules.smartcache.SmartcacheDatamodule.setup"]], "test_dataloader() (cyto_dl.datamodules.smartcache.smartcachedatamodule method)": [[33, "cyto_dl.datamodules.smartcache.SmartcacheDatamodule.test_dataloader"]], "train_dataloader() (cyto_dl.datamodules.smartcache.smartcachedatamodule method)": [[33, "cyto_dl.datamodules.smartcache.SmartcacheDatamodule.train_dataloader"]], "val_dataloader() (cyto_dl.datamodules.smartcache.smartcachedatamodule method)": [[33, "cyto_dl.datamodules.smartcache.SmartcacheDatamodule.val_dataloader"]], "cifar10 (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.CIFAR10"]], "cifar100 (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.CIFAR100"]], "clevrclassification (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.CLEVRClassification"]], "caltech101 (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.Caltech101"]], "caltech256 (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.Caltech256"]], "celeba (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.CelebA"]], "cityscapes (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.Cityscapes"]], "cococaptions (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.CocoCaptions"]], "cocodetection (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.CocoDetection"]], "country211 (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.Country211"]], "dtd (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.DTD"]], "datasetfolder (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.DatasetFolder"]], "emnist (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.EMNIST"]], "eurosat (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.EuroSAT"]], "fer2013 (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.FER2013"]], "fgvcaircraft (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.FGVCAircraft"]], "fakedata (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.FakeData"]], "fashionmnist (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.FashionMNIST"]], "flickr30k (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.Flickr30k"]], "flickr8k (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.Flickr8k"]], "flowers102 (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.Flowers102"]], "food101 (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.Food101"]], "gtsrb (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.GTSRB"]], "hmdb51 (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.HMDB51"]], "inaturalist (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.INaturalist"]], "imagefolder (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.ImageFolder"]], "imagenet (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.ImageNet"]], "kmnist (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.KMNIST"]], "kinetics (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.Kinetics"]], "kitti (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.Kitti"]], "lfwpairs (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.LFWPairs"]], "lfwpeople (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.LFWPeople"]], "lsun (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.LSUN"]], "lsunclass (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.LSUNClass"]], "mnist (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.MNIST"]], "movingmnist (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.MovingMNIST"]], "omniglot (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.Omniglot"]], "oxfordiiitpet (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.OxfordIIITPet"]], "pcam (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.PCAM"]], "phototour (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.PhotoTour"]], "places365 (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.Places365"]], "qmnist (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.QMNIST"]], "renderedsst2 (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.RenderedSST2"]], "sbdataset (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.SBDataset"]], "sbu (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.SBU"]], "semeion (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.SEMEION"]], "stl10 (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.STL10"]], "sun397 (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.SUN397"]], "svhn (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.SVHN"]], "stanfordcars (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.StanfordCars"]], "ucf101 (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.UCF101"]], "usps (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.USPS"]], "vocdetection (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.VOCDetection"]], "vocsegmentation (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.VOCSegmentation"]], "visiondataset (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.VisionDataset"]], "widerface (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.WIDERFace"]], "cyto_dl.datamodules.torchvision": [[34, "module-cyto_dl.datamodules.torchvision"]], "cyto_dl.eval": [[35, "module-cyto_dl.eval"]], "main() (in module cyto_dl.eval)": [[35, "cyto_dl.eval.main"]], "cyto_dl.image": [[36, "module-cyto_dl.image"]], "cyto_dl.image.io": [[37, "module-cyto_dl.image.io"]], "aicsimageloaderd (class in cyto_dl.image.io.aicsimage_loader)": [[38, "cyto_dl.image.io.aicsimage_loader.AICSImageLoaderd"]], "cyto_dl.image.io.aicsimage_loader": [[38, "module-cyto_dl.image.io.aicsimage_loader"]], "split_args() (cyto_dl.image.io.aicsimage_loader.aicsimageloaderd method)": [[38, "cyto_dl.image.io.aicsimage_loader.AICSImageLoaderd.split_args"]], "monaibioreader (class in cyto_dl.image.io.monai_bio_reader)": [[39, "cyto_dl.image.io.monai_bio_reader.MonaiBioReader"]], "cyto_dl.image.io.monai_bio_reader": [[39, "module-cyto_dl.image.io.monai_bio_reader"]], "get_data() (cyto_dl.image.io.monai_bio_reader.monaibioreader method)": [[39, "cyto_dl.image.io.monai_bio_reader.MonaiBioReader.get_data"]], "read() (cyto_dl.image.io.monai_bio_reader.monaibioreader method)": [[39, "cyto_dl.image.io.monai_bio_reader.MonaiBioReader.read"]], "verify_suffix() (cyto_dl.image.io.monai_bio_reader.monaibioreader method)": [[39, "cyto_dl.image.io.monai_bio_reader.MonaiBioReader.verify_suffix"]], "readnumpyfile (class in cyto_dl.image.io.numpy_reader)": [[40, "cyto_dl.image.io.numpy_reader.ReadNumpyFile"]], "cyto_dl.image.io.numpy_reader": [[40, "module-cyto_dl.image.io.numpy_reader"]], "omezarrreader (class in cyto_dl.image.io.ome_zarr_reader)": [[41, "cyto_dl.image.io.ome_zarr_reader.OmeZarrReader"]], "cyto_dl.image.io.ome_zarr_reader": [[41, "module-cyto_dl.image.io.ome_zarr_reader"]], "get_data() (cyto_dl.image.io.ome_zarr_reader.omezarrreader method)": [[41, "cyto_dl.image.io.ome_zarr_reader.OmeZarrReader.get_data"]], "read() (cyto_dl.image.io.ome_zarr_reader.omezarrreader method)": [[41, "cyto_dl.image.io.ome_zarr_reader.OmeZarrReader.read"]], "verify_suffix() (cyto_dl.image.io.ome_zarr_reader.omezarrreader method)": [[41, "cyto_dl.image.io.ome_zarr_reader.OmeZarrReader.verify_suffix"]], "polygonloaderd (class in cyto_dl.image.io.polygon_loader)": [[42, "cyto_dl.image.io.polygon_loader.PolygonLoaderd"]], "cyto_dl.image.io.polygon_loader": [[42, "module-cyto_dl.image.io.polygon_loader"]], "skimagereader (class in cyto_dl.image.io.skimage_reader)": [[43, "cyto_dl.image.io.skimage_reader.SkimageReader"]], "cyto_dl.image.io.skimage_reader": [[43, "module-cyto_dl.image.io.skimage_reader"]], "get_data() (cyto_dl.image.io.skimage_reader.skimagereader method)": [[43, "cyto_dl.image.io.skimage_reader.SkimageReader.get_data"]], "read() (cyto_dl.image.io.skimage_reader.skimagereader method)": [[43, "cyto_dl.image.io.skimage_reader.SkimageReader.read"]], "verify_suffix() (cyto_dl.image.io.skimage_reader.skimagereader method)": [[43, "cyto_dl.image.io.skimage_reader.SkimageReader.verify_suffix"]], "cyto_dl.image.transforms": [[44, "module-cyto_dl.image.transforms"]], "addmeta (class in cyto_dl.image.transforms.add_meta)": [[45, "cyto_dl.image.transforms.add_meta.AddMeta"]], "cyto_dl.image.transforms.add_meta": [[45, "module-cyto_dl.image.transforms.add_meta"]], "brightsampler (class in cyto_dl.image.transforms.bright_sampler)": [[46, "cyto_dl.image.transforms.bright_sampler.BrightSampler"]], "cyto_dl.image.transforms.bright_sampler": [[46, "module-cyto_dl.image.transforms.bright_sampler"]], "clip (class in cyto_dl.image.transforms.clip)": [[47, "cyto_dl.image.transforms.clip.Clip"]], "clipd (class in cyto_dl.image.transforms.clip)": [[47, "cyto_dl.image.transforms.clip.Clipd"]], "cyto_dl.image.transforms.clip": [[47, "module-cyto_dl.image.transforms.clip"]], "contrastadjust (class in cyto_dl.image.transforms.contrastadjust)": [[48, "cyto_dl.image.transforms.contrastadjust.ContrastAdjust"]], "contrastadjustd (class in cyto_dl.image.transforms.contrastadjust)": [[48, "cyto_dl.image.transforms.contrastadjust.ContrastAdjustd"]], "cyto_dl.image.transforms.contrastadjust": [[48, "module-cyto_dl.image.transforms.contrastadjust"]], "jepamaskgenerator (class in cyto_dl.image.transforms.generate_jepa_masks)": [[49, "cyto_dl.image.transforms.generate_jepa_masks.JEPAMaskGenerator"]], "cyto_dl.image.transforms.generate_jepa_masks": [[49, "module-cyto_dl.image.transforms.generate_jepa_masks"]], "remove_excess_pixels() (cyto_dl.image.transforms.generate_jepa_masks.jepamaskgenerator method)": [[49, "cyto_dl.image.transforms.generate_jepa_masks.JEPAMaskGenerator.remove_excess_pixels"]], "merged (class in cyto_dl.image.transforms.merge)": [[50, "cyto_dl.image.transforms.merge.Merged"]], "cyto_dl.image.transforms.merge": [[50, "module-cyto_dl.image.transforms.merge"]], "randommultiscalecropd (class in cyto_dl.image.transforms.multiscale_cropper)": [[51, "cyto_dl.image.transforms.multiscale_cropper.RandomMultiScaleCropd"]], "cyto_dl.image.transforms.multiscale_cropper": [[51, "module-cyto_dl.image.transforms.multiscale_cropper"]], "generate_slices() (cyto_dl.image.transforms.multiscale_cropper.randommultiscalecropd method)": [[51, "cyto_dl.image.transforms.multiscale_cropper.RandomMultiScaleCropd.generate_slices"]], "padzd (class in cyto_dl.image.transforms.pad)": [[52, "cyto_dl.image.transforms.pad.PadZd"]], "cyto_dl.image.transforms.pad": [[52, "module-cyto_dl.image.transforms.pad"]], "maxprojectd (class in cyto_dl.image.transforms.project)": [[53, "cyto_dl.image.transforms.project.MaxProjectd"]], "cyto_dl.image.transforms.project": [[53, "module-cyto_dl.image.transforms.project"]], "save (class in cyto_dl.image.transforms.save)": [[55, "cyto_dl.image.transforms.save.Save"]], "saved (class in cyto_dl.image.transforms.save)": [[55, "cyto_dl.image.transforms.save.Saved"]], "cyto_dl.image.transforms.save": [[55, "module-cyto_dl.image.transforms.save"]], "cropresize (class in cyto_dl.image.transforms.track_transforms)": [[57, "cyto_dl.image.transforms.track_transforms.CropResize"]], "generatetracklabels (class in cyto_dl.image.transforms.track_transforms)": [[57, "cyto_dl.image.transforms.track_transforms.GenerateTrackLabels"]], "perchannel (class in cyto_dl.image.transforms.track_transforms)": [[57, "cyto_dl.image.transforms.track_transforms.PerChannel"]], "splittrackd (class in cyto_dl.image.transforms.track_transforms)": [[57, "cyto_dl.image.transforms.track_transforms.SplitTrackd"]], "cyto_dl.image.transforms.track_transforms": [[57, "module-cyto_dl.image.transforms.track_transforms"]], "cyto_dl.loggers": [[58, "module-cyto_dl.loggers"]], "mlflowlogger (class in cyto_dl.loggers.mlflow)": [[59, "cyto_dl.loggers.mlflow.MLFlowLogger"]], "after_save_checkpoint() (cyto_dl.loggers.mlflow.mlflowlogger method)": [[59, "cyto_dl.loggers.mlflow.MLFlowLogger.after_save_checkpoint"]], "cyto_dl.loggers.mlflow": [[59, "module-cyto_dl.loggers.mlflow"]], "log_hyperparams() (cyto_dl.loggers.mlflow.mlflowlogger method)": [[59, "cyto_dl.loggers.mlflow.MLFlowLogger.log_hyperparams"]], "log_metrics() (cyto_dl.loggers.mlflow.mlflowlogger method)": [[59, "cyto_dl.loggers.mlflow.MLFlowLogger.log_metrics"]], "cyto_dl.models": [[60, "module-cyto_dl.models"]], "basemodel (class in cyto_dl.models.base_model)": [[61, "cyto_dl.models.base_model.BaseModel"]], "basemodelmeta (class in cyto_dl.models.base_model)": [[61, "cyto_dl.models.base_model.BaseModelMeta"]], "compute_metrics() (cyto_dl.models.base_model.basemodel method)": [[61, "cyto_dl.models.base_model.BaseModel.compute_metrics"]], "configure_optimizers() (cyto_dl.models.base_model.basemodel method)": [[61, "cyto_dl.models.base_model.BaseModel.configure_optimizers"]], "cyto_dl.models.base_model": [[61, "module-cyto_dl.models.base_model"]], "forward() (cyto_dl.models.base_model.basemodel method)": [[61, "cyto_dl.models.base_model.BaseModel.forward"]], "model_step() (cyto_dl.models.base_model.basemodel method)": [[61, "cyto_dl.models.base_model.BaseModel.model_step"]], "on_train_start() (cyto_dl.models.base_model.basemodel method)": [[61, "cyto_dl.models.base_model.BaseModel.on_train_start"]], "parse_batch() (cyto_dl.models.base_model.basemodel method)": [[61, "cyto_dl.models.base_model.BaseModel.parse_batch"]], "predict_step() (cyto_dl.models.base_model.basemodel method)": [[61, "cyto_dl.models.base_model.BaseModel.predict_step"]], "test_step() (cyto_dl.models.base_model.basemodel method)": [[61, "cyto_dl.models.base_model.BaseModel.test_step"]], "training_step() (cyto_dl.models.base_model.basemodel method)": [[61, "cyto_dl.models.base_model.BaseModel.training_step"]], "validation_step() (cyto_dl.models.base_model.basemodel method)": [[61, "cyto_dl.models.base_model.BaseModel.validation_step"]], "basicmodel (class in cyto_dl.models.basic_model)": [[62, "cyto_dl.models.basic_model.BasicModel"]], "cyto_dl.models.basic_model": [[62, "module-cyto_dl.models.basic_model"]], "forward() (cyto_dl.models.basic_model.basicmodel method)": [[62, "cyto_dl.models.basic_model.BasicModel.forward"]], "model_step() (cyto_dl.models.basic_model.basicmodel method)": [[62, "cyto_dl.models.basic_model.BasicModel.model_step"]], "cyto_dl.models.classification": [[63, "module-cyto_dl.models.classification"]], "classification (class in cyto_dl.models.classification.classification)": [[64, "cyto_dl.models.classification.classification.Classification"]], "cyto_dl.models.classification.classification": [[64, "module-cyto_dl.models.classification.classification"]], "forward() (cyto_dl.models.classification.classification.classification method)": [[64, "cyto_dl.models.classification.classification.Classification.forward"]], "model_step() (cyto_dl.models.classification.classification.classification method)": [[64, "cyto_dl.models.classification.classification.Classification.model_step"]], "predict_step() (cyto_dl.models.classification.classification.classification method)": [[64, "cyto_dl.models.classification.classification.Classification.predict_step"]], "save_images() (cyto_dl.models.classification.classification.classification method)": [[64, "cyto_dl.models.classification.classification.Classification.save_images"]], "should_save_image() (cyto_dl.models.classification.classification.classification method)": [[64, "cyto_dl.models.classification.classification.Classification.should_save_image"]], "timepointclassification (class in cyto_dl.models.classification.timepoint_classification)": [[65, "cyto_dl.models.classification.timepoint_classification.TimepointClassification"]], "cyto_dl.models.classification.timepoint_classification": [[65, "module-cyto_dl.models.classification.timepoint_classification"]], "predict_step() (cyto_dl.models.classification.timepoint_classification.timepointclassification method)": [[65, "cyto_dl.models.classification.timepoint_classification.TimepointClassification.predict_step"]], "cyto_dl.models.contrastive": [[66, "module-cyto_dl.models.contrastive"]], "contrastive (class in cyto_dl.models.contrastive.contrastive)": [[67, "cyto_dl.models.contrastive.contrastive.Contrastive"]], "cyto_dl.models.contrastive.contrastive": [[67, "module-cyto_dl.models.contrastive.contrastive"]], "forward() (cyto_dl.models.contrastive.contrastive.contrastive method)": [[67, "cyto_dl.models.contrastive.contrastive.Contrastive.forward"]], "model_step() (cyto_dl.models.contrastive.contrastive.contrastive method)": [[67, "cyto_dl.models.contrastive.contrastive.Contrastive.model_step"]], "plot_classes() (cyto_dl.models.contrastive.contrastive.contrastive method)": [[67, "cyto_dl.models.contrastive.contrastive.Contrastive.plot_classes"]], "plot_neighbors() (cyto_dl.models.contrastive.contrastive.contrastive method)": [[67, "cyto_dl.models.contrastive.contrastive.Contrastive.plot_neighbors"]], "predict_step() (cyto_dl.models.contrastive.contrastive.contrastive method)": [[67, "cyto_dl.models.contrastive.contrastive.Contrastive.predict_step"]], "cyto_dl.models.im2im": [[72, "module-cyto_dl.models.im2im"]], "gan (class in cyto_dl.models.im2im.gan)": [[73, "cyto_dl.models.im2im.gan.GAN"]], "configure_optimizers() (cyto_dl.models.im2im.gan.gan method)": [[73, "cyto_dl.models.im2im.gan.GAN.configure_optimizers"]], "cyto_dl.models.im2im.gan": [[73, "module-cyto_dl.models.im2im.gan"]], "model_step() (cyto_dl.models.im2im.gan.gan method)": [[73, "cyto_dl.models.im2im.gan.GAN.model_step"]], "predict_step() (cyto_dl.models.im2im.gan.gan method)": [[73, "cyto_dl.models.im2im.gan.GAN.predict_step"]], "multitaskim2im (class in cyto_dl.models.im2im.multi_task)": [[74, "cyto_dl.models.im2im.multi_task.MultiTaskIm2Im"]], "configure_optimizers() (cyto_dl.models.im2im.multi_task.multitaskim2im method)": [[74, "cyto_dl.models.im2im.multi_task.MultiTaskIm2Im.configure_optimizers"]], "cyto_dl.models.im2im.multi_task": [[74, "module-cyto_dl.models.im2im.multi_task"]], "forward() (cyto_dl.models.im2im.multi_task.multitaskim2im method)": [[74, "cyto_dl.models.im2im.multi_task.MultiTaskIm2Im.forward"]], "get_n_postprocess_image() (cyto_dl.models.im2im.multi_task.multitaskim2im method)": [[74, "cyto_dl.models.im2im.multi_task.MultiTaskIm2Im.get_n_postprocess_image"]], "get_per_head() (cyto_dl.models.im2im.multi_task.multitaskim2im method)": [[74, "cyto_dl.models.im2im.multi_task.MultiTaskIm2Im.get_per_head"]], "model_step() (cyto_dl.models.im2im.multi_task.multitaskim2im method)": [[74, "cyto_dl.models.im2im.multi_task.MultiTaskIm2Im.model_step"]], "predict_step() (cyto_dl.models.im2im.multi_task.multitaskim2im method)": [[74, "cyto_dl.models.im2im.multi_task.MultiTaskIm2Im.predict_step"]], "run_forward() (cyto_dl.models.im2im.multi_task.multitaskim2im method)": [[74, "cyto_dl.models.im2im.multi_task.MultiTaskIm2Im.run_forward"]], "test_step() (cyto_dl.models.im2im.multi_task.multitaskim2im method)": [[74, "cyto_dl.models.im2im.multi_task.MultiTaskIm2Im.test_step"]], "training_step() (cyto_dl.models.im2im.multi_task.multitaskim2im method)": [[74, "cyto_dl.models.im2im.multi_task.MultiTaskIm2Im.training_step"]], "validation_step() (cyto_dl.models.im2im.multi_task.multitaskim2im method)": [[74, "cyto_dl.models.im2im.multi_task.MultiTaskIm2Im.validation_step"]], "cyto_dl.models.im2im.utils": [[75, "module-cyto_dl.models.im2im.utils"]], "instancesegcluster (class in cyto_dl.models.im2im.utils.instance_seg)": [[76, "cyto_dl.models.im2im.utils.instance_seg.InstanceSegCluster"]], "instancesegloss (class in cyto_dl.models.im2im.utils.instance_seg)": [[76, "cyto_dl.models.im2im.utils.instance_seg.InstanceSegLoss"]], "instancesegpreprocessd (class in cyto_dl.models.im2im.utils.instance_seg)": [[76, "cyto_dl.models.im2im.utils.instance_seg.InstanceSegPreprocessd"]], "instancesegrandflipd (class in cyto_dl.models.im2im.utils.instance_seg)": [[76, "cyto_dl.models.im2im.utils.instance_seg.InstanceSegRandFlipd"]], "cluster_object() (cyto_dl.models.im2im.utils.instance_seg.instancesegcluster method)": [[76, "cyto_dl.models.im2im.utils.instance_seg.InstanceSegCluster.cluster_object"]], "cyto_dl.models.im2im.utils.instance_seg": [[76, "module-cyto_dl.models.im2im.utils.instance_seg"]], "embed_from_skel() (cyto_dl.models.im2im.utils.instance_seg.instancesegpreprocessd method)": [[76, "cyto_dl.models.im2im.utils.instance_seg.InstanceSegPreprocessd.embed_from_skel"]], "kd_clustering() (cyto_dl.models.im2im.utils.instance_seg.instancesegcluster method)": [[76, "cyto_dl.models.im2im.utils.instance_seg.InstanceSegCluster.kd_clustering"]], "keep_largest_cc() (cyto_dl.models.im2im.utils.instance_seg.instancesegpreprocessd method)": [[76, "cyto_dl.models.im2im.utils.instance_seg.InstanceSegPreprocessd.keep_largest_cc"]], "label_2d() (cyto_dl.models.im2im.utils.instance_seg.instancesegpreprocessd method)": [[76, "cyto_dl.models.im2im.utils.instance_seg.InstanceSegPreprocessd.label_2d"]], "pad_slice() (in module cyto_dl.models.im2im.utils.instance_seg)": [[76, "cyto_dl.models.im2im.utils.instance_seg.pad_slice"]], "remove_small_skeletons() (cyto_dl.models.im2im.utils.instance_seg.instancesegcluster method)": [[76, "cyto_dl.models.im2im.utils.instance_seg.InstanceSegCluster.remove_small_skeletons"]], "shrink() (cyto_dl.models.im2im.utils.instance_seg.instancesegpreprocessd method)": [[76, "cyto_dl.models.im2im.utils.instance_seg.InstanceSegPreprocessd.shrink"]], "skeleton_tall() (cyto_dl.models.im2im.utils.instance_seg.instancesegpreprocessd method)": [[76, "cyto_dl.models.im2im.utils.instance_seg.InstanceSegPreprocessd.skeleton_tall"]], "smooth_embedding() (cyto_dl.models.im2im.utils.instance_seg.instancesegpreprocessd method)": [[76, "cyto_dl.models.im2im.utils.instance_seg.InstanceSegPreprocessd.smooth_embedding"]], "topology_preserving_thinning() (cyto_dl.models.im2im.utils.instance_seg.instancesegpreprocessd method)": [[76, "cyto_dl.models.im2im.utils.instance_seg.InstanceSegPreprocessd.topology_preserving_thinning"]], "noiseannealer (class in cyto_dl.models.im2im.utils.noise_annealer)": [[77, "cyto_dl.models.im2im.utils.noise_annealer.NoiseAnnealer"]], "cyto_dl.models.im2im.utils.noise_annealer": [[77, "module-cyto_dl.models.im2im.utils.noise_annealer"]], "update_noise() (cyto_dl.models.im2im.utils.noise_annealer.noiseannealer method)": [[77, "cyto_dl.models.im2im.utils.noise_annealer.NoiseAnnealer.update_noise"]], "cyto_dl.models.im2im.utils.postprocessing": [[78, "module-cyto_dl.models.im2im.utils.postprocessing"]], "actthreshlabel (class in cyto_dl.models.im2im.utils.postprocessing.act_thresh_label)": [[79, "cyto_dl.models.im2im.utils.postprocessing.act_thresh_label.ActThreshLabel"]], "cyto_dl.models.im2im.utils.postprocessing.act_thresh_label": [[79, "module-cyto_dl.models.im2im.utils.postprocessing.act_thresh_label"]], "autothreshold (class in cyto_dl.models.im2im.utils.postprocessing.auto_thresh)": [[80, "cyto_dl.models.im2im.utils.postprocessing.auto_thresh.AutoThreshold"]], "cyto_dl.models.im2im.utils.postprocessing.auto_thresh": [[80, "module-cyto_dl.models.im2im.utils.postprocessing.auto_thresh"]], "dicttoim (class in cyto_dl.models.im2im.utils.postprocessing.dict_to_im)": [[81, "cyto_dl.models.im2im.utils.postprocessing.dict_to_im.DictToIm"]], "cyto_dl.models.im2im.utils.postprocessing.dict_to_im": [[81, "module-cyto_dl.models.im2im.utils.postprocessing.dict_to_im"]], "detach() (in module cyto_dl.models.im2im.utils.postprocessing.dict_to_im)": [[81, "cyto_dl.models.im2im.utils.postprocessing.dict_to_im.detach"]], "cyto_dl.models.jepa": [[82, "module-cyto_dl.models.jepa"]], "ijepa (class in cyto_dl.models.jepa.ijepa)": [[83, "cyto_dl.models.jepa.ijepa.IJEPA"]], "cyto_dl.models.jepa.ijepa": [[83, "module-cyto_dl.models.jepa.ijepa"]], "model_step() (cyto_dl.models.jepa.ijepa.ijepa method)": [[83, "cyto_dl.models.jepa.ijepa.IJEPA.model_step"]], "iwm (class in cyto_dl.models.jepa.iwm)": [[84, "cyto_dl.models.jepa.iwm.IWM"]], "cyto_dl.models.jepa.iwm": [[84, "module-cyto_dl.models.jepa.iwm"]], "extract_embeddings() (cyto_dl.models.jepa.iwm.iwm method)": [[84, "cyto_dl.models.jepa.iwm.IWM.extract_embeddings"]], "get_predict_masks() (cyto_dl.models.jepa.iwm.iwm method)": [[84, "cyto_dl.models.jepa.iwm.IWM.get_predict_masks"]], "model_step() (cyto_dl.models.jepa.iwm.iwm method)": [[84, "cyto_dl.models.jepa.iwm.IWM.model_step"]], "predict_step() (cyto_dl.models.jepa.iwm.iwm method)": [[84, "cyto_dl.models.jepa.iwm.IWM.predict_step"]], "test_step() (cyto_dl.models.jepa.iwm.iwm method)": [[84, "cyto_dl.models.jepa.iwm.IWM.test_step"]], "jepabase (class in cyto_dl.models.jepa.jepa_base)": [[85, "cyto_dl.models.jepa.jepa_base.JEPABase"]], "configure_optimizers() (cyto_dl.models.jepa.jepa_base.jepabase method)": [[85, "cyto_dl.models.jepa.jepa_base.JEPABase.configure_optimizers"]], "cyto_dl.models.jepa.jepa_base": [[85, "module-cyto_dl.models.jepa.jepa_base"]], "forward() (cyto_dl.models.jepa.jepa_base.jepabase method)": [[85, "cyto_dl.models.jepa.jepa_base.JEPABase.forward"]], "get_context_embeddings() (cyto_dl.models.jepa.jepa_base.jepabase method)": [[85, "cyto_dl.models.jepa.jepa_base.JEPABase.get_context_embeddings"]], "get_mask() (cyto_dl.models.jepa.jepa_base.jepabase method)": [[85, "cyto_dl.models.jepa.jepa_base.JEPABase.get_mask"]], "get_target_embeddings() (cyto_dl.models.jepa.jepa_base.jepabase method)": [[85, "cyto_dl.models.jepa.jepa_base.JEPABase.get_target_embeddings"]], "model_step() (cyto_dl.models.jepa.jepa_base.jepabase method)": [[85, "cyto_dl.models.jepa.jepa_base.JEPABase.model_step"]], "predict_step() (cyto_dl.models.jepa.jepa_base.jepabase method)": [[85, "cyto_dl.models.jepa.jepa_base.JEPABase.predict_step"]], "remove_first_dim() (cyto_dl.models.jepa.jepa_base.jepabase method)": [[85, "cyto_dl.models.jepa.jepa_base.JEPABase.remove_first_dim"]], "update_teacher() (cyto_dl.models.jepa.jepa_base.jepabase method)": [[85, "cyto_dl.models.jepa.jepa_base.JEPABase.update_teacher"]], "cyto_dl.models.utils": [[86, "module-cyto_dl.models.utils"]], "cyto_dl.models.utils.mlflow": [[87, "module-cyto_dl.models.utils.mlflow"]], "get_config() (in module cyto_dl.models.utils.mlflow)": [[87, "cyto_dl.models.utils.mlflow.get_config"]], "load_model_from_checkpoint() (in module cyto_dl.models.utils.mlflow)": [[87, "cyto_dl.models.utils.mlflow.load_model_from_checkpoint"]], "cyto_dl.models.utils.utils": [[88, "module-cyto_dl.models.utils.utils"]], "find_indices() (in module cyto_dl.models.utils.utils)": [[88, "cyto_dl.models.utils.utils.find_indices"]], "basevae (class in cyto_dl.models.vae.base_vae)": [[90, "cyto_dl.models.vae.base_vae.BaseVAE"]], "calculate_elbo() (cyto_dl.models.vae.base_vae.basevae method)": [[90, "cyto_dl.models.vae.base_vae.BaseVAE.calculate_elbo"]], "calculate_rcl() (cyto_dl.models.vae.base_vae.basevae method)": [[90, "cyto_dl.models.vae.base_vae.BaseVAE.calculate_rcl"]], "calculate_rcl_dict() (cyto_dl.models.vae.base_vae.basevae method)": [[90, "cyto_dl.models.vae.base_vae.BaseVAE.calculate_rcl_dict"]], "cyto_dl.models.vae.base_vae": [[90, "module-cyto_dl.models.vae.base_vae"]], "decode() (cyto_dl.models.vae.base_vae.basevae method)": [[90, "cyto_dl.models.vae.base_vae.BaseVAE.decode"]], "encode() (cyto_dl.models.vae.base_vae.basevae method)": [[90, "cyto_dl.models.vae.base_vae.BaseVAE.encode"]], "forward() (cyto_dl.models.vae.base_vae.basevae method)": [[90, "cyto_dl.models.vae.base_vae.BaseVAE.forward"]], "model_step() (cyto_dl.models.vae.base_vae.basevae method)": [[90, "cyto_dl.models.vae.base_vae.BaseVAE.model_step"]], "sample_z() (cyto_dl.models.vae.base_vae.basevae method)": [[90, "cyto_dl.models.vae.base_vae.BaseVAE.sample_z"]], "cyto_dl.models.vae.priors": [[99, "module-cyto_dl.models.vae.priors"]], "prior (class in cyto_dl.models.vae.priors.abstract_prior)": [[100, "cyto_dl.models.vae.priors.abstract_prior.Prior"]], "cyto_dl.models.vae.priors.abstract_prior": [[100, "module-cyto_dl.models.vae.priors.abstract_prior"]], "forward() (cyto_dl.models.vae.priors.abstract_prior.prior method)": [[100, "cyto_dl.models.vae.priors.abstract_prior.Prior.forward"]], "kl_divergence() (cyto_dl.models.vae.priors.abstract_prior.prior method)": [[100, "cyto_dl.models.vae.priors.abstract_prior.Prior.kl_divergence"]], "param_size (cyto_dl.models.vae.priors.abstract_prior.prior property)": [[100, "cyto_dl.models.vae.priors.abstract_prior.Prior.param_size"]], "sample() (cyto_dl.models.vae.priors.abstract_prior.prior method)": [[100, "cyto_dl.models.vae.priors.abstract_prior.Prior.sample"]], "diagonalgaussianprior (class in cyto_dl.models.vae.priors.gaussian)": [[101, "cyto_dl.models.vae.priors.gaussian.DiagonalGaussianPrior"]], "isotropicgaussianprior (class in cyto_dl.models.vae.priors.gaussian)": [[101, "cyto_dl.models.vae.priors.gaussian.IsotropicGaussianPrior"]], "compute_tc_penalty() (in module cyto_dl.models.vae.priors.gaussian)": [[101, "cyto_dl.models.vae.priors.gaussian.compute_tc_penalty"]], "cyto_dl.models.vae.priors.gaussian": [[101, "module-cyto_dl.models.vae.priors.gaussian"]], "forward() (cyto_dl.models.vae.priors.gaussian.diagonalgaussianprior method)": [[101, "cyto_dl.models.vae.priors.gaussian.DiagonalGaussianPrior.forward"]], "forward() (cyto_dl.models.vae.priors.gaussian.isotropicgaussianprior method)": [[101, "cyto_dl.models.vae.priors.gaussian.IsotropicGaussianPrior.forward"]], "kl_divergence() (cyto_dl.models.vae.priors.gaussian.diagonalgaussianprior class method)": [[101, "cyto_dl.models.vae.priors.gaussian.DiagonalGaussianPrior.kl_divergence"]], "kl_divergence() (cyto_dl.models.vae.priors.gaussian.isotropicgaussianprior class method)": [[101, "cyto_dl.models.vae.priors.gaussian.IsotropicGaussianPrior.kl_divergence"]], "param_size (cyto_dl.models.vae.priors.gaussian.diagonalgaussianprior property)": [[101, "cyto_dl.models.vae.priors.gaussian.DiagonalGaussianPrior.param_size"]], "param_size (cyto_dl.models.vae.priors.gaussian.isotropicgaussianprior property)": [[101, "cyto_dl.models.vae.priors.gaussian.IsotropicGaussianPrior.param_size"]], "sample() (cyto_dl.models.vae.priors.gaussian.isotropicgaussianprior class method)": [[101, "cyto_dl.models.vae.priors.gaussian.IsotropicGaussianPrior.sample"]], "identityprior (class in cyto_dl.models.vae.priors.identity_prior)": [[102, "cyto_dl.models.vae.priors.identity_prior.IdentityPrior"]], "cyto_dl.models.vae.priors.identity_prior": [[102, "module-cyto_dl.models.vae.priors.identity_prior"]], "forward() (cyto_dl.models.vae.priors.identity_prior.identityprior method)": [[102, "cyto_dl.models.vae.priors.identity_prior.IdentityPrior.forward"]], "param_size (cyto_dl.models.vae.priors.identity_prior.identityprior property)": [[102, "cyto_dl.models.vae.priors.identity_prior.IdentityPrior.param_size"]], "jointprior (class in cyto_dl.models.vae.priors.joint_prior)": [[103, "cyto_dl.models.vae.priors.joint_prior.JointPrior"]], "cyto_dl.models.vae.priors.joint_prior": [[103, "module-cyto_dl.models.vae.priors.joint_prior"]], "forward() (cyto_dl.models.vae.priors.joint_prior.jointprior method)": [[103, "cyto_dl.models.vae.priors.joint_prior.JointPrior.forward"]], "kl_divergence() (cyto_dl.models.vae.priors.joint_prior.jointprior method)": [[103, "cyto_dl.models.vae.priors.joint_prior.JointPrior.kl_divergence"]], "param_size (cyto_dl.models.vae.priors.joint_prior.jointprior property)": [[103, "cyto_dl.models.vae.priors.joint_prior.JointPrior.param_size"]], "sample() (cyto_dl.models.vae.priors.joint_prior.jointprior method)": [[103, "cyto_dl.models.vae.priors.joint_prior.JointPrior.sample"]], "cyto_dl.nn": [[105, "module-cyto_dl.nn"]], "cyto_dl.nn.discriminators": [[106, "module-cyto_dl.nn.discriminators"]], "multiscalediscriminator (class in cyto_dl.nn.discriminators.multi_scale_discriminator)": [[107, "cyto_dl.nn.discriminators.multi_scale_discriminator.MultiScaleDiscriminator"]], "cyto_dl.nn.discriminators.multi_scale_discriminator": [[107, "module-cyto_dl.nn.discriminators.multi_scale_discriminator"]], "forward() (cyto_dl.nn.discriminators.multi_scale_discriminator.multiscalediscriminator method)": [[107, "cyto_dl.nn.discriminators.multi_scale_discriminator.MultiScaleDiscriminator.forward"]], "nlayerdiscriminator (class in cyto_dl.nn.discriminators.n_layer_discriminator)": [[108, "cyto_dl.nn.discriminators.n_layer_discriminator.NLayerDiscriminator"]], "cyto_dl.nn.discriminators.n_layer_discriminator": [[108, "module-cyto_dl.nn.discriminators.n_layer_discriminator"]], "forward() (cyto_dl.nn.discriminators.n_layer_discriminator.nlayerdiscriminator method)": [[108, "cyto_dl.nn.discriminators.n_layer_discriminator.NLayerDiscriminator.forward"]], "cyto_dl.nn.head": [[109, "module-cyto_dl.nn.head"]], "basehead (class in cyto_dl.nn.head.base_head)": [[110, "cyto_dl.nn.head.base_head.BaseHead"]], "cyto_dl.nn.head.base_head": [[110, "module-cyto_dl.nn.head.base_head"]], "forward() (cyto_dl.nn.head.base_head.basehead method)": [[110, "cyto_dl.nn.head.base_head.BaseHead.forward"]], "generate_io_map() (cyto_dl.nn.head.base_head.basehead method)": [[110, "cyto_dl.nn.head.base_head.BaseHead.generate_io_map"]], "run_head() (cyto_dl.nn.head.base_head.basehead method)": [[110, "cyto_dl.nn.head.base_head.BaseHead.run_head"]], "update_params() (cyto_dl.nn.head.base_head.basehead method)": [[110, "cyto_dl.nn.head.base_head.BaseHead.update_params"]], "ganhead (class in cyto_dl.nn.head.gan_head)": [[111, "cyto_dl.nn.head.gan_head.GANHead"]], "cyto_dl.nn.head.gan_head": [[111, "module-cyto_dl.nn.head.gan_head"]], "forward() (cyto_dl.nn.head.gan_head.ganhead method)": [[111, "cyto_dl.nn.head.gan_head.GANHead.forward"]], "run_head() (cyto_dl.nn.head.gan_head.ganhead method)": [[111, "cyto_dl.nn.head.gan_head.GANHead.run_head"]], "ganhead_resize (class in cyto_dl.nn.head.gan_head_superres)": [[112, "cyto_dl.nn.head.gan_head_superres.GANHead_resize"]], "cyto_dl.nn.head.gan_head_superres": [[112, "module-cyto_dl.nn.head.gan_head_superres"]], "forward() (cyto_dl.nn.head.gan_head_superres.ganhead_resize method)": [[112, "cyto_dl.nn.head.gan_head_superres.GANHead_resize.forward"]], "maehead (class in cyto_dl.nn.head.mae_head)": [[113, "cyto_dl.nn.head.mae_head.MAEHead"]], "cyto_dl.nn.head.mae_head": [[113, "module-cyto_dl.nn.head.mae_head"]], "run_head() (cyto_dl.nn.head.mae_head.maehead method)": [[113, "cyto_dl.nn.head.mae_head.MAEHead.run_head"]], "maskhead (class in cyto_dl.nn.head.mask_head)": [[114, "cyto_dl.nn.head.mask_head.MaskHead"]], "cyto_dl.nn.head.mask_head": [[114, "module-cyto_dl.nn.head.mask_head"]], "run_head() (cyto_dl.nn.head.mask_head.maskhead method)": [[114, "cyto_dl.nn.head.mask_head.MaskHead.run_head"]], "resblockshead (class in cyto_dl.nn.head.res_blocks_head)": [[115, "cyto_dl.nn.head.res_blocks_head.ResBlocksHead"]], "cyto_dl.nn.head.res_blocks_head": [[115, "module-cyto_dl.nn.head.res_blocks_head"]], "forward() (cyto_dl.nn.head.res_blocks_head.resblockshead method)": [[115, "cyto_dl.nn.head.res_blocks_head.ResBlocksHead.forward"]], "vicreghead (class in cyto_dl.nn.head.vic_reg)": [[116, "cyto_dl.nn.head.vic_reg.VICRegHead"]], "cyto_dl.nn.head.vic_reg": [[116, "module-cyto_dl.nn.head.vic_reg"]], "run_head() (cyto_dl.nn.head.vic_reg.vicreghead method)": [[116, "cyto_dl.nn.head.vic_reg.VICRegHead.run_head"]], "hrskip (class in cyto_dl.nn.hr_skip)": [[117, "cyto_dl.nn.hr_skip.HRSkip"]], "cyto_dl.nn.hr_skip": [[117, "module-cyto_dl.nn.hr_skip"]], "forward() (cyto_dl.nn.hr_skip.hrskip method)": [[117, "cyto_dl.nn.hr_skip.HRSkip.forward"]], "cyto_dl.nn.losses": [[118, "module-cyto_dl.nn.losses"]], "adversarialloss (class in cyto_dl.nn.losses.adversarial_loss)": [[119, "cyto_dl.nn.losses.adversarial_loss.AdversarialLoss"]], "cyto_dl.nn.losses.adversarial_loss": [[119, "module-cyto_dl.nn.losses.adversarial_loss"]], "forward() (cyto_dl.nn.losses.adversarial_loss.adversarialloss method)": [[119, "cyto_dl.nn.losses.adversarial_loss.AdversarialLoss.forward"]], "chamferloss (class in cyto_dl.nn.losses.chamfer_loss)": [[120, "cyto_dl.nn.losses.chamfer_loss.ChamferLoss"]], "batch_pairwise_dist() (cyto_dl.nn.losses.chamfer_loss.chamferloss method)": [[120, "cyto_dl.nn.losses.chamfer_loss.ChamferLoss.batch_pairwise_dist"]], "batch_pairwise_dist2() (cyto_dl.nn.losses.chamfer_loss.chamferloss method)": [[120, "cyto_dl.nn.losses.chamfer_loss.ChamferLoss.batch_pairwise_dist2"]], "cyto_dl.nn.losses.chamfer_loss": [[120, "module-cyto_dl.nn.losses.chamfer_loss"]], "forward() (cyto_dl.nn.losses.chamfer_loss.chamferloss method)": [[120, "cyto_dl.nn.losses.chamfer_loss.ChamferLoss.forward"]], "cblogloss (class in cyto_dl.nn.losses.continuous_bernoulli)": [[121, "cyto_dl.nn.losses.continuous_bernoulli.CBLogLoss"]], "cyto_dl.nn.losses.continuous_bernoulli": [[121, "module-cyto_dl.nn.losses.continuous_bernoulli"]], "forward() (cyto_dl.nn.losses.continuous_bernoulli.cblogloss method)": [[121, "cyto_dl.nn.losses.continuous_bernoulli.CBLogLoss.forward"]], "cosineloss (class in cyto_dl.nn.losses.cosine_loss)": [[122, "cyto_dl.nn.losses.cosine_loss.CosineLoss"]], "cyto_dl.nn.losses.cosine_loss": [[122, "module-cyto_dl.nn.losses.cosine_loss"]], "forward() (cyto_dl.nn.losses.cosine_loss.cosineloss method)": [[122, "cyto_dl.nn.losses.cosine_loss.CosineLoss.forward"]], "ganloss (class in cyto_dl.nn.losses.gan_loss)": [[123, "cyto_dl.nn.losses.gan_loss.GANLoss"]], "pix2pixhd (class in cyto_dl.nn.losses.gan_loss)": [[123, "cyto_dl.nn.losses.gan_loss.Pix2PixHD"]], "cyto_dl.nn.losses.gan_loss": [[123, "module-cyto_dl.nn.losses.gan_loss"]], "get_feature_matching_loss() (cyto_dl.nn.losses.gan_loss.pix2pixhd method)": [[123, "cyto_dl.nn.losses.gan_loss.Pix2PixHD.get_feature_matching_loss"]], "get_gan_loss() (cyto_dl.nn.losses.gan_loss.pix2pixhd method)": [[123, "cyto_dl.nn.losses.gan_loss.Pix2PixHD.get_gan_loss"]], "get_target_tensor() (cyto_dl.nn.losses.gan_loss.ganloss method)": [[123, "cyto_dl.nn.losses.gan_loss.GANLoss.get_target_tensor"]], "gaussiannllloss (class in cyto_dl.nn.losses.gaussian_nll_loss)": [[124, "cyto_dl.nn.losses.gaussian_nll_loss.GaussianNLLLoss"]], "cyto_dl.nn.losses.gaussian_nll_loss": [[124, "module-cyto_dl.nn.losses.gaussian_nll_loss"]], "forward() (cyto_dl.nn.losses.gaussian_nll_loss.gaussiannllloss method)": [[124, "cyto_dl.nn.losses.gaussian_nll_loss.GaussianNLLLoss.forward"]], "cmap_loss (class in cyto_dl.nn.losses.loss_wrapper)": [[126, "cyto_dl.nn.losses.loss_wrapper.CMAP_loss"]], "losswrapper (class in cyto_dl.nn.losses.loss_wrapper)": [[126, "cyto_dl.nn.losses.loss_wrapper.LossWrapper"]], "cyto_dl.nn.losses.loss_wrapper": [[126, "module-cyto_dl.nn.losses.loss_wrapper"]], "thresholdloss (class in cyto_dl.nn.losses.threshold_loss)": [[128, "cyto_dl.nn.losses.threshold_loss.ThresholdLoss"]], "cyto_dl.nn.losses.threshold_loss": [[128, "module-cyto_dl.nn.losses.threshold_loss"]], "vicregloss (class in cyto_dl.nn.losses.vic_reg)": [[129, "cyto_dl.nn.losses.vic_reg.VICRegLoss"]], "cyto_dl.nn.losses.vic_reg": [[129, "module-cyto_dl.nn.losses.vic_reg"]], "forward() (cyto_dl.nn.losses.vic_reg.vicregloss method)": [[129, "cyto_dl.nn.losses.vic_reg.VICRegLoss.forward"]], "off_diagonal() (in module cyto_dl.nn.losses.vic_reg)": [[129, "cyto_dl.nn.losses.vic_reg.off_diagonal"]], "weibulllogloss (class in cyto_dl.nn.losses.weibull)": [[130, "cyto_dl.nn.losses.weibull.WeibullLogLoss"]], "cyto_dl.nn.losses.weibull": [[130, "module-cyto_dl.nn.losses.weibull"]], "forward() (cyto_dl.nn.losses.weibull.weibulllogloss method)": [[130, "cyto_dl.nn.losses.weibull.WeibullLogLoss.forward"]], "weibull_log_probs() (in module cyto_dl.nn.losses.weibull)": [[130, "cyto_dl.nn.losses.weibull.weibull_log_probs"]], "weightedmseloss (class in cyto_dl.nn.losses.weighted_mse_loss)": [[131, "cyto_dl.nn.losses.weighted_mse_loss.WeightedMSELoss"]], "cyto_dl.nn.losses.weighted_mse_loss": [[131, "module-cyto_dl.nn.losses.weighted_mse_loss"]], "forward() (cyto_dl.nn.losses.weighted_mse_loss.weightedmseloss method)": [[131, "cyto_dl.nn.losses.weighted_mse_loss.WeightedMSELoss.forward"]], "mlp (class in cyto_dl.nn.mlp)": [[132, "cyto_dl.nn.mlp.MLP"]], "cyto_dl.nn.mlp": [[132, "module-cyto_dl.nn.mlp"]], "forward() (cyto_dl.nn.mlp.mlp method)": [[132, "cyto_dl.nn.mlp.MLP.forward"]], "residualunit (class in cyto_dl.nn.res_unit)": [[138, "cyto_dl.nn.res_unit.ResidualUnit"]], "cyto_dl.nn.res_unit": [[138, "module-cyto_dl.nn.res_unit"]], "forward() (cyto_dl.nn.res_unit.residualunit method)": [[138, "cyto_dl.nn.res_unit.ResidualUnit.forward"]], "convpoolrelu (class in cyto_dl.nn.spatial_transformer)": [[139, "cyto_dl.nn.spatial_transformer.ConvPoolReLU"]], "stn (class in cyto_dl.nn.spatial_transformer)": [[139, "cyto_dl.nn.spatial_transformer.STN"]], "cyto_dl.nn.spatial_transformer": [[139, "module-cyto_dl.nn.spatial_transformer"]], "forward() (cyto_dl.nn.spatial_transformer.convpoolrelu method)": [[139, "cyto_dl.nn.spatial_transformer.ConvPoolReLU.forward"]], "forward() (cyto_dl.nn.spatial_transformer.stn method)": [[139, "cyto_dl.nn.spatial_transformer.STN.forward"]], "toggle() (cyto_dl.nn.spatial_transformer.stn method)": [[139, "cyto_dl.nn.spatial_transformer.STN.toggle"]], "trackclassifier (class in cyto_dl.nn.track_sequence_predictor)": [[140, "cyto_dl.nn.track_sequence_predictor.TrackClassifier"]], "cyto_dl.nn.track_sequence_predictor": [[140, "module-cyto_dl.nn.track_sequence_predictor"]], "forward() (cyto_dl.nn.track_sequence_predictor.trackclassifier method)": [[140, "cyto_dl.nn.track_sequence_predictor.TrackClassifier.forward"]], "positionalencoding1d() (in module cyto_dl.nn.track_sequence_predictor)": [[140, "cyto_dl.nn.track_sequence_predictor.positionalencoding1d"]], "cyto_dl.nn.vits": [[141, "module-cyto_dl.nn.vits"]], "cyto_dl.nn.vits.blocks": [[142, "module-cyto_dl.nn.vits.blocks"]], "crossattention (class in cyto_dl.nn.vits.blocks.cross_attention)": [[143, "cyto_dl.nn.vits.blocks.cross_attention.CrossAttention"]], "crossattentionblock (class in cyto_dl.nn.vits.blocks.cross_attention)": [[143, "cyto_dl.nn.vits.blocks.cross_attention.CrossAttentionBlock"]], "crossselfblock (class in cyto_dl.nn.vits.blocks.cross_attention)": [[143, "cyto_dl.nn.vits.blocks.cross_attention.CrossSelfBlock"]], "mlp (class in cyto_dl.nn.vits.blocks.cross_attention)": [[143, "cyto_dl.nn.vits.blocks.cross_attention.Mlp"]], "cyto_dl.nn.vits.blocks.cross_attention": [[143, "module-cyto_dl.nn.vits.blocks.cross_attention"]], "forward() (cyto_dl.nn.vits.blocks.cross_attention.crossattention method)": [[143, "cyto_dl.nn.vits.blocks.cross_attention.CrossAttention.forward"]], "forward() (cyto_dl.nn.vits.blocks.cross_attention.crossattentionblock method)": [[143, "cyto_dl.nn.vits.blocks.cross_attention.CrossAttentionBlock.forward"]], "forward() (cyto_dl.nn.vits.blocks.cross_attention.crossselfblock method)": [[143, "cyto_dl.nn.vits.blocks.cross_attention.CrossSelfBlock.forward"]], "forward() (cyto_dl.nn.vits.blocks.cross_attention.mlp method)": [[143, "cyto_dl.nn.vits.blocks.cross_attention.Mlp.forward"]], "intermediateweigher (class in cyto_dl.nn.vits.blocks.intermediate_weigher)": [[144, "cyto_dl.nn.vits.blocks.intermediate_weigher.IntermediateWeigher"]], "cyto_dl.nn.vits.blocks.intermediate_weigher": [[144, "module-cyto_dl.nn.vits.blocks.intermediate_weigher"]], "forward() (cyto_dl.nn.vits.blocks.intermediate_weigher.intermediateweigher method)": [[144, "cyto_dl.nn.vits.blocks.intermediate_weigher.IntermediateWeigher.forward"]], "patchify (class in cyto_dl.nn.vits.blocks.patchify)": [[145, "cyto_dl.nn.vits.blocks.patchify.Patchify"]], "cyto_dl.nn.vits.blocks.patchify": [[145, "module-cyto_dl.nn.vits.blocks.patchify"]], "forward() (cyto_dl.nn.vits.blocks.patchify.patchify method)": [[145, "cyto_dl.nn.vits.blocks.patchify.Patchify.forward"]], "get_mask() (cyto_dl.nn.vits.blocks.patchify.patchify method)": [[145, "cyto_dl.nn.vits.blocks.patchify.Patchify.get_mask"]], "random_indexes() (in module cyto_dl.nn.vits.blocks.patchify)": [[145, "cyto_dl.nn.vits.blocks.patchify.random_indexes"]], "crossmae_decoder (class in cyto_dl.nn.vits.cross_mae)": [[146, "cyto_dl.nn.vits.cross_mae.CrossMAE_Decoder"]], "cyto_dl.nn.vits.cross_mae": [[146, "module-cyto_dl.nn.vits.cross_mae"]], "forward() (cyto_dl.nn.vits.cross_mae.crossmae_decoder method)": [[146, "cyto_dl.nn.vits.cross_mae.CrossMAE_Decoder.forward"]], "init_weight() (cyto_dl.nn.vits.cross_mae.crossmae_decoder method)": [[146, "cyto_dl.nn.vits.cross_mae.CrossMAE_Decoder.init_weight"]], "iwmpredictor (class in cyto_dl.nn.vits.jepa)": [[147, "cyto_dl.nn.vits.jepa.IWMPredictor"]], "jepaencoder (class in cyto_dl.nn.vits.jepa)": [[147, "cyto_dl.nn.vits.jepa.JEPAEncoder"]], "jepapredictor (class in cyto_dl.nn.vits.jepa)": [[147, "cyto_dl.nn.vits.jepa.JEPAPredictor"]], "cyto_dl.nn.vits.jepa": [[147, "module-cyto_dl.nn.vits.jepa"]], "forward() (cyto_dl.nn.vits.jepa.iwmpredictor method)": [[147, "cyto_dl.nn.vits.jepa.IWMPredictor.forward"]], "forward() (cyto_dl.nn.vits.jepa.jepaencoder method)": [[147, "cyto_dl.nn.vits.jepa.JEPAEncoder.forward"]], "forward() (cyto_dl.nn.vits.jepa.jepapredictor method)": [[147, "cyto_dl.nn.vits.jepa.JEPAPredictor.forward"]], "init_weight() (cyto_dl.nn.vits.jepa.jepapredictor method)": [[147, "cyto_dl.nn.vits.jepa.JEPAPredictor.init_weight"]], "predict_target_features() (cyto_dl.nn.vits.jepa.jepapredictor method)": [[147, "cyto_dl.nn.vits.jepa.JEPAPredictor.predict_target_features"]], "mae_decoder (class in cyto_dl.nn.vits.mae)": [[148, "cyto_dl.nn.vits.mae.MAE_Decoder"]], "mae_encoder (class in cyto_dl.nn.vits.mae)": [[148, "cyto_dl.nn.vits.mae.MAE_Encoder"]], "mae_vit (class in cyto_dl.nn.vits.mae)": [[148, "cyto_dl.nn.vits.mae.MAE_ViT"]], "cyto_dl.nn.vits.mae": [[148, "module-cyto_dl.nn.vits.mae"]], "forward() (cyto_dl.nn.vits.mae.mae_decoder method)": [[148, "cyto_dl.nn.vits.mae.MAE_Decoder.forward"]], "forward() (cyto_dl.nn.vits.mae.mae_encoder method)": [[148, "cyto_dl.nn.vits.mae.MAE_Encoder.forward"]], "forward() (cyto_dl.nn.vits.mae.mae_vit method)": [[148, "cyto_dl.nn.vits.mae.MAE_ViT.forward"]], "init_weight() (cyto_dl.nn.vits.mae.mae_decoder method)": [[148, "cyto_dl.nn.vits.mae.MAE_Decoder.init_weight"]], "init_weight() (cyto_dl.nn.vits.mae.mae_encoder method)": [[148, "cyto_dl.nn.vits.mae.MAE_Encoder.init_weight"]], "encodedskip (class in cyto_dl.nn.vits.seg)": [[149, "cyto_dl.nn.vits.seg.EncodedSkip"]], "seg_vit (class in cyto_dl.nn.vits.seg)": [[149, "cyto_dl.nn.vits.seg.Seg_ViT"]], "superresdecoder (class in cyto_dl.nn.vits.seg)": [[149, "cyto_dl.nn.vits.seg.SuperresDecoder"]], "cyto_dl.nn.vits.seg": [[149, "module-cyto_dl.nn.vits.seg"]], "forward() (cyto_dl.nn.vits.seg.encodedskip method)": [[149, "cyto_dl.nn.vits.seg.EncodedSkip.forward"]], "forward() (cyto_dl.nn.vits.seg.seg_vit method)": [[149, "cyto_dl.nn.vits.seg.Seg_ViT.forward"]], "forward() (cyto_dl.nn.vits.seg.superresdecoder method)": [[149, "cyto_dl.nn.vits.seg.SuperresDecoder.forward"]], "cyto_dl.nn.vits.utils": [[150, "module-cyto_dl.nn.vits.utils"]], "get_positional_embedding() (in module cyto_dl.nn.vits.utils)": [[150, "cyto_dl.nn.vits.utils.get_positional_embedding"]], "take_indexes() (in module cyto_dl.nn.vits.utils)": [[150, "cyto_dl.nn.vits.utils.take_indexes"]], "cyto_dl.point_cloud": [[151, "module-cyto_dl.point_cloud"]], "cyto_dl.train": [[154, "module-cyto_dl.train"]], "main() (in module cyto_dl.train)": [[154, "cyto_dl.train.main"]], "cyto_dl.utils": [[155, "module-cyto_dl.utils"]], "create_dataloader() (in module cyto_dl.utils.array)": [[156, "cyto_dl.utils.array.create_dataloader"]], "cyto_dl.utils.array": [[156, "module-cyto_dl.utils.array"]], "extract_array_predictions() (in module cyto_dl.utils.array)": [[156, "cyto_dl.utils.array.extract_array_predictions"]], "cyto_dl.utils.checkpoint": [[157, "module-cyto_dl.utils.checkpoint"]], "load_checkpoint() (in module cyto_dl.utils.checkpoint)": [[157, "cyto_dl.utils.checkpoint.load_checkpoint"]], "cyto_dl.utils.config": [[158, "module-cyto_dl.utils.config"]], "is_config() (in module cyto_dl.utils.config)": [[158, "cyto_dl.utils.config.is_config"]], "kv_to_dict() (in module cyto_dl.utils.config)": [[158, "cyto_dl.utils.config.kv_to_dict"]], "remove_aux_key() (in module cyto_dl.utils.config)": [[158, "cyto_dl.utils.config.remove_aux_key"]], "cyto_dl.utils.download_test_data": [[159, "module-cyto_dl.utils.download_test_data"]], "delete_test_data() (in module cyto_dl.utils.download_test_data)": [[159, "cyto_dl.utils.download_test_data.delete_test_data"]], "download_test_data() (in module cyto_dl.utils.download_test_data)": [[159, "cyto_dl.utils.download_test_data.download_test_data"]], "parse_s3_path() (in module cyto_dl.utils.download_test_data)": [[159, "cyto_dl.utils.download_test_data.parse_s3_path"]], "setup_paths() (in module cyto_dl.utils.download_test_data)": [[159, "cyto_dl.utils.download_test_data.setup_paths"]], "dummydatamodule (class in cyto_dl.utils.dummy_dataset)": [[160, "cyto_dl.utils.dummy_dataset.DummyDatamodule"]], "dummydataset (class in cyto_dl.utils.dummy_dataset)": [[160, "cyto_dl.utils.dummy_dataset.DummyDataset"]], "cyto_dl.utils.dummy_dataset": [[160, "module-cyto_dl.utils.dummy_dataset"]], "generate_img() (cyto_dl.utils.dummy_dataset.dummydataset method)": [[160, "cyto_dl.utils.dummy_dataset.DummyDataset.generate_img"]], "get_dataloader() (cyto_dl.utils.dummy_dataset.dummydatamodule method)": [[160, "cyto_dl.utils.dummy_dataset.DummyDatamodule.get_dataloader"]], "predict_dataloader() (cyto_dl.utils.dummy_dataset.dummydatamodule method)": [[160, "cyto_dl.utils.dummy_dataset.DummyDatamodule.predict_dataloader"]], "test_dataloader() (cyto_dl.utils.dummy_dataset.dummydatamodule method)": [[160, "cyto_dl.utils.dummy_dataset.DummyDatamodule.test_dataloader"]], "train_dataloader() (cyto_dl.utils.dummy_dataset.dummydatamodule method)": [[160, "cyto_dl.utils.dummy_dataset.DummyDatamodule.train_dataloader"]], "val_dataloader() (cyto_dl.utils.dummy_dataset.dummydatamodule method)": [[160, "cyto_dl.utils.dummy_dataset.DummyDatamodule.val_dataloader"]], "cyto_dl.utils.pylogger": [[161, "module-cyto_dl.utils.pylogger"]], "get_pylogger() (in module cyto_dl.utils.pylogger)": [[161, "cyto_dl.utils.pylogger.get_pylogger"]], "cyto_dl.utils.rich_utils": [[162, "module-cyto_dl.utils.rich_utils"]], "enforce_tags() (in module cyto_dl.utils.rich_utils)": [[162, "cyto_dl.utils.rich_utils.enforce_tags"]], "print_config_tree() (in module cyto_dl.utils.rich_utils)": [[162, "cyto_dl.utils.rich_utils.print_config_tree"]], "rotationmodule (class in cyto_dl.utils.rotation)": [[163, "cyto_dl.utils.rotation.RotationModule"]], "compute_rotation_matrix() (cyto_dl.utils.rotation.rotationmodule method)": [[163, "cyto_dl.utils.rotation.RotationModule.compute_rotation_matrix"]], "cyto_dl.utils.rotation": [[163, "module-cyto_dl.utils.rotation"]], "get_rotation_matrix_so2() (in module cyto_dl.utils.rotation)": [[163, "cyto_dl.utils.rotation.get_rotation_matrix_so2"]], "get_rotation_matrix_so3() (in module cyto_dl.utils.rotation)": [[163, "cyto_dl.utils.rotation.get_rotation_matrix_so3"]], "close_loggers() (in module cyto_dl.utils.template_utils)": [[167, "cyto_dl.utils.template_utils.close_loggers"]], "cyto_dl.utils.template_utils": [[167, "module-cyto_dl.utils.template_utils"]], "extras() (in module cyto_dl.utils.template_utils)": [[167, "cyto_dl.utils.template_utils.extras"]], "get_metric_value() (in module cyto_dl.utils.template_utils)": [[167, "cyto_dl.utils.template_utils.get_metric_value"]], "instantiate_callbacks() (in module cyto_dl.utils.template_utils)": [[167, "cyto_dl.utils.template_utils.instantiate_callbacks"]], "instantiate_loggers() (in module cyto_dl.utils.template_utils)": [[167, "cyto_dl.utils.template_utils.instantiate_loggers"]], "log_hyperparameters() (in module cyto_dl.utils.template_utils)": [[167, "cyto_dl.utils.template_utils.log_hyperparameters"]], "save_file() (in module cyto_dl.utils.template_utils)": [[167, "cyto_dl.utils.template_utils.save_file"]], "task_wrapper() (in module cyto_dl.utils.template_utils)": [[167, "cyto_dl.utils.template_utils.task_wrapper"]]}})
\ No newline at end of file
+Search.setIndex({"docnames": ["api", "cyto_dl", "cyto_dl.api", "cyto_dl.api.cyto_dl_model", "cyto_dl.api.cyto_dl_model.cyto_dl_base_model", "cyto_dl.api.cyto_dl_model.segmentation_plugin_model", "cyto_dl.api.data", "cyto_dl.api.model", "cyto_dl.callbacks", "cyto_dl.callbacks.callback_utils", "cyto_dl.callbacks.csv_saver", "cyto_dl.callbacks.image_saver", "cyto_dl.callbacks.latent_walk", "cyto_dl.callbacks.layer_freeze", "cyto_dl.callbacks.model_utils", "cyto_dl.callbacks.outlier_detection", "cyto_dl.compile", "cyto_dl.dataframe", "cyto_dl.dataframe.readers", "cyto_dl.dataframe.transforms", "cyto_dl.dataframe.transforms.filter", "cyto_dl.dataframe.transforms.group_cols", "cyto_dl.dataframe.transforms.misc", "cyto_dl.dataframe.transforms.split", "cyto_dl.datamodules", "cyto_dl.datamodules.array", "cyto_dl.datamodules.data_dict", "cyto_dl.datamodules.dataframe", "cyto_dl.datamodules.dataframe.dataframe_datamodule", "cyto_dl.datamodules.dataframe.grouped_dataframe_datamodule", "cyto_dl.datamodules.dataframe.utils", "cyto_dl.datamodules.folder", "cyto_dl.datamodules.multidim_image", "cyto_dl.datamodules.smartcache", "cyto_dl.datamodules.torchvision", "cyto_dl.eval", "cyto_dl.image", "cyto_dl.image.io", "cyto_dl.image.io.aicsimage_loader", "cyto_dl.image.io.monai_bio_reader", "cyto_dl.image.io.numpy_reader", "cyto_dl.image.io.ome_zarr_reader", "cyto_dl.image.io.polygon_loader", "cyto_dl.image.io.skimage_reader", "cyto_dl.image.transforms", "cyto_dl.image.transforms.add_meta", "cyto_dl.image.transforms.bright_sampler", "cyto_dl.image.transforms.clip", "cyto_dl.image.transforms.contrastadjust", "cyto_dl.image.transforms.generate_jepa_masks", "cyto_dl.image.transforms.merge", "cyto_dl.image.transforms.multiscale_cropper", "cyto_dl.image.transforms.pad", "cyto_dl.image.transforms.project", "cyto_dl.image.transforms.rotation_mask_transform", "cyto_dl.image.transforms.save", "cyto_dl.image.transforms.so2_random_rotation", "cyto_dl.image.transforms.track_transforms", "cyto_dl.loggers", "cyto_dl.loggers.mlflow", "cyto_dl.models", "cyto_dl.models.base_model", "cyto_dl.models.basic_model", "cyto_dl.models.classification", "cyto_dl.models.classification.classification", "cyto_dl.models.classification.timepoint_classification", "cyto_dl.models.contrastive", "cyto_dl.models.contrastive.contrastive", "cyto_dl.models.handlers", "cyto_dl.models.handlers.base_handler", "cyto_dl.models.handlers.image_handler", "cyto_dl.models.handlers.load_image_patch", "cyto_dl.models.im2im", "cyto_dl.models.im2im.gan", "cyto_dl.models.im2im.multi_task", "cyto_dl.models.im2im.utils", "cyto_dl.models.im2im.utils.instance_seg", "cyto_dl.models.im2im.utils.noise_annealer", "cyto_dl.models.im2im.utils.postprocessing", "cyto_dl.models.im2im.utils.postprocessing.act_thresh_label", "cyto_dl.models.im2im.utils.postprocessing.auto_thresh", "cyto_dl.models.im2im.utils.postprocessing.dict_to_im", "cyto_dl.models.jepa", "cyto_dl.models.jepa.ijepa", "cyto_dl.models.jepa.iwm", "cyto_dl.models.jepa.jepa_base", "cyto_dl.models.utils", "cyto_dl.models.utils.mlflow", "cyto_dl.models.utils.utils", "cyto_dl.models.vae", "cyto_dl.models.vae.base_vae", "cyto_dl.models.vae.image_canon_vae", "cyto_dl.models.vae.image_encoder", "cyto_dl.models.vae.image_vae", "cyto_dl.models.vae.latent_loss_vae", "cyto_dl.models.vae.o2_spharm_vae", "cyto_dl.models.vae.o2_spharm_vae.o2_spharm_encoder", "cyto_dl.models.vae.o2_spharm_vae.o2_spharm_vae", "cyto_dl.models.vae.point_cloud_vae", "cyto_dl.models.vae.priors", "cyto_dl.models.vae.priors.abstract_prior", "cyto_dl.models.vae.priors.gaussian", "cyto_dl.models.vae.priors.identity_prior", "cyto_dl.models.vae.priors.joint_prior", "cyto_dl.models.vae.tabular_vae", "cyto_dl.nn", "cyto_dl.nn.discriminators", "cyto_dl.nn.discriminators.multi_scale_discriminator", "cyto_dl.nn.discriminators.n_layer_discriminator", "cyto_dl.nn.head", "cyto_dl.nn.head.base_head", "cyto_dl.nn.head.gan_head", "cyto_dl.nn.head.gan_head_superres", "cyto_dl.nn.head.mae_head", "cyto_dl.nn.head.mask_head", "cyto_dl.nn.head.res_blocks_head", "cyto_dl.nn.head.vic_reg", "cyto_dl.nn.hr_skip", "cyto_dl.nn.losses", "cyto_dl.nn.losses.adversarial_loss", "cyto_dl.nn.losses.chamfer_loss", "cyto_dl.nn.losses.continuous_bernoulli", "cyto_dl.nn.losses.cosine_loss", "cyto_dl.nn.losses.gan_loss", "cyto_dl.nn.losses.gaussian_nll_loss", "cyto_dl.nn.losses.geomloss", "cyto_dl.nn.losses.loss_wrapper", "cyto_dl.nn.losses.spharm_loss", "cyto_dl.nn.losses.threshold_loss", "cyto_dl.nn.losses.vic_reg", "cyto_dl.nn.losses.weibull", "cyto_dl.nn.losses.weighted_mse_loss", "cyto_dl.nn.mlp", "cyto_dl.nn.point_cloud", "cyto_dl.nn.point_cloud.dgcnn", "cyto_dl.nn.point_cloud.folding_net", "cyto_dl.nn.point_cloud.graph_functions", "cyto_dl.nn.point_cloud.vnn", "cyto_dl.nn.res_unit", "cyto_dl.nn.spatial_transformer", "cyto_dl.nn.track_sequence_predictor", "cyto_dl.nn.vits", "cyto_dl.nn.vits.blocks", "cyto_dl.nn.vits.blocks.cross_attention", "cyto_dl.nn.vits.blocks.intermediate_weigher", "cyto_dl.nn.vits.blocks.masked_unit_attention", "cyto_dl.nn.vits.blocks.patchify", "cyto_dl.nn.vits.blocks.patchify.patchify", "cyto_dl.nn.vits.blocks.patchify.patchify_base", "cyto_dl.nn.vits.blocks.patchify.patchify_hiera", "cyto_dl.nn.vits.decoder", "cyto_dl.nn.vits.encoder", "cyto_dl.nn.vits.mae", "cyto_dl.nn.vits.predictor", "cyto_dl.nn.vits.seg", "cyto_dl.nn.vits.utils", "cyto_dl.point_cloud", "cyto_dl.point_cloud.io", "cyto_dl.point_cloud.io.read_pcloud", "cyto_dl.train", "cyto_dl.utils", "cyto_dl.utils.array", "cyto_dl.utils.checkpoint", "cyto_dl.utils.config", "cyto_dl.utils.download_test_data", "cyto_dl.utils.dummy_dataset", "cyto_dl.utils.pylogger", "cyto_dl.utils.rich_utils", "cyto_dl.utils.rotation", "cyto_dl.utils.spharm", "cyto_dl.utils.spharm.reconstruction", "cyto_dl.utils.spharm.rotation", "cyto_dl.utils.template_utils", "index", "installation", "modules", "overview", "using_examples"], "filenames": ["api.rst", "cyto_dl.rst", "cyto_dl.api.rst", "cyto_dl.api.cyto_dl_model.rst", "cyto_dl.api.cyto_dl_model.cyto_dl_base_model.rst", "cyto_dl.api.cyto_dl_model.segmentation_plugin_model.rst", "cyto_dl.api.data.rst", "cyto_dl.api.model.rst", "cyto_dl.callbacks.rst", "cyto_dl.callbacks.callback_utils.rst", "cyto_dl.callbacks.csv_saver.rst", "cyto_dl.callbacks.image_saver.rst", "cyto_dl.callbacks.latent_walk.rst", "cyto_dl.callbacks.layer_freeze.rst", "cyto_dl.callbacks.model_utils.rst", "cyto_dl.callbacks.outlier_detection.rst", "cyto_dl.compile.rst", "cyto_dl.dataframe.rst", "cyto_dl.dataframe.readers.rst", "cyto_dl.dataframe.transforms.rst", "cyto_dl.dataframe.transforms.filter.rst", "cyto_dl.dataframe.transforms.group_cols.rst", "cyto_dl.dataframe.transforms.misc.rst", "cyto_dl.dataframe.transforms.split.rst", "cyto_dl.datamodules.rst", "cyto_dl.datamodules.array.rst", "cyto_dl.datamodules.data_dict.rst", "cyto_dl.datamodules.dataframe.rst", "cyto_dl.datamodules.dataframe.dataframe_datamodule.rst", "cyto_dl.datamodules.dataframe.grouped_dataframe_datamodule.rst", "cyto_dl.datamodules.dataframe.utils.rst", "cyto_dl.datamodules.folder.rst", "cyto_dl.datamodules.multidim_image.rst", "cyto_dl.datamodules.smartcache.rst", "cyto_dl.datamodules.torchvision.rst", "cyto_dl.eval.rst", "cyto_dl.image.rst", "cyto_dl.image.io.rst", "cyto_dl.image.io.aicsimage_loader.rst", "cyto_dl.image.io.monai_bio_reader.rst", "cyto_dl.image.io.numpy_reader.rst", "cyto_dl.image.io.ome_zarr_reader.rst", "cyto_dl.image.io.polygon_loader.rst", "cyto_dl.image.io.skimage_reader.rst", "cyto_dl.image.transforms.rst", "cyto_dl.image.transforms.add_meta.rst", "cyto_dl.image.transforms.bright_sampler.rst", "cyto_dl.image.transforms.clip.rst", "cyto_dl.image.transforms.contrastadjust.rst", "cyto_dl.image.transforms.generate_jepa_masks.rst", "cyto_dl.image.transforms.merge.rst", "cyto_dl.image.transforms.multiscale_cropper.rst", "cyto_dl.image.transforms.pad.rst", "cyto_dl.image.transforms.project.rst", "cyto_dl.image.transforms.rotation_mask_transform.rst", "cyto_dl.image.transforms.save.rst", "cyto_dl.image.transforms.so2_random_rotation.rst", "cyto_dl.image.transforms.track_transforms.rst", "cyto_dl.loggers.rst", "cyto_dl.loggers.mlflow.rst", "cyto_dl.models.rst", "cyto_dl.models.base_model.rst", "cyto_dl.models.basic_model.rst", "cyto_dl.models.classification.rst", "cyto_dl.models.classification.classification.rst", "cyto_dl.models.classification.timepoint_classification.rst", "cyto_dl.models.contrastive.rst", "cyto_dl.models.contrastive.contrastive.rst", "cyto_dl.models.handlers.rst", "cyto_dl.models.handlers.base_handler.rst", "cyto_dl.models.handlers.image_handler.rst", "cyto_dl.models.handlers.load_image_patch.rst", "cyto_dl.models.im2im.rst", "cyto_dl.models.im2im.gan.rst", "cyto_dl.models.im2im.multi_task.rst", "cyto_dl.models.im2im.utils.rst", "cyto_dl.models.im2im.utils.instance_seg.rst", "cyto_dl.models.im2im.utils.noise_annealer.rst", "cyto_dl.models.im2im.utils.postprocessing.rst", "cyto_dl.models.im2im.utils.postprocessing.act_thresh_label.rst", "cyto_dl.models.im2im.utils.postprocessing.auto_thresh.rst", "cyto_dl.models.im2im.utils.postprocessing.dict_to_im.rst", "cyto_dl.models.jepa.rst", "cyto_dl.models.jepa.ijepa.rst", "cyto_dl.models.jepa.iwm.rst", "cyto_dl.models.jepa.jepa_base.rst", "cyto_dl.models.utils.rst", "cyto_dl.models.utils.mlflow.rst", "cyto_dl.models.utils.utils.rst", "cyto_dl.models.vae.rst", "cyto_dl.models.vae.base_vae.rst", "cyto_dl.models.vae.image_canon_vae.rst", "cyto_dl.models.vae.image_encoder.rst", "cyto_dl.models.vae.image_vae.rst", "cyto_dl.models.vae.latent_loss_vae.rst", "cyto_dl.models.vae.o2_spharm_vae.rst", "cyto_dl.models.vae.o2_spharm_vae.o2_spharm_encoder.rst", "cyto_dl.models.vae.o2_spharm_vae.o2_spharm_vae.rst", "cyto_dl.models.vae.point_cloud_vae.rst", "cyto_dl.models.vae.priors.rst", "cyto_dl.models.vae.priors.abstract_prior.rst", "cyto_dl.models.vae.priors.gaussian.rst", "cyto_dl.models.vae.priors.identity_prior.rst", "cyto_dl.models.vae.priors.joint_prior.rst", "cyto_dl.models.vae.tabular_vae.rst", "cyto_dl.nn.rst", "cyto_dl.nn.discriminators.rst", "cyto_dl.nn.discriminators.multi_scale_discriminator.rst", "cyto_dl.nn.discriminators.n_layer_discriminator.rst", "cyto_dl.nn.head.rst", "cyto_dl.nn.head.base_head.rst", "cyto_dl.nn.head.gan_head.rst", "cyto_dl.nn.head.gan_head_superres.rst", "cyto_dl.nn.head.mae_head.rst", "cyto_dl.nn.head.mask_head.rst", "cyto_dl.nn.head.res_blocks_head.rst", "cyto_dl.nn.head.vic_reg.rst", "cyto_dl.nn.hr_skip.rst", "cyto_dl.nn.losses.rst", "cyto_dl.nn.losses.adversarial_loss.rst", "cyto_dl.nn.losses.chamfer_loss.rst", "cyto_dl.nn.losses.continuous_bernoulli.rst", "cyto_dl.nn.losses.cosine_loss.rst", "cyto_dl.nn.losses.gan_loss.rst", "cyto_dl.nn.losses.gaussian_nll_loss.rst", "cyto_dl.nn.losses.geomloss.rst", "cyto_dl.nn.losses.loss_wrapper.rst", "cyto_dl.nn.losses.spharm_loss.rst", "cyto_dl.nn.losses.threshold_loss.rst", "cyto_dl.nn.losses.vic_reg.rst", "cyto_dl.nn.losses.weibull.rst", "cyto_dl.nn.losses.weighted_mse_loss.rst", "cyto_dl.nn.mlp.rst", "cyto_dl.nn.point_cloud.rst", "cyto_dl.nn.point_cloud.dgcnn.rst", "cyto_dl.nn.point_cloud.folding_net.rst", "cyto_dl.nn.point_cloud.graph_functions.rst", "cyto_dl.nn.point_cloud.vnn.rst", "cyto_dl.nn.res_unit.rst", "cyto_dl.nn.spatial_transformer.rst", "cyto_dl.nn.track_sequence_predictor.rst", "cyto_dl.nn.vits.rst", "cyto_dl.nn.vits.blocks.rst", "cyto_dl.nn.vits.blocks.cross_attention.rst", "cyto_dl.nn.vits.blocks.intermediate_weigher.rst", "cyto_dl.nn.vits.blocks.masked_unit_attention.rst", "cyto_dl.nn.vits.blocks.patchify.rst", "cyto_dl.nn.vits.blocks.patchify.patchify.rst", "cyto_dl.nn.vits.blocks.patchify.patchify_base.rst", "cyto_dl.nn.vits.blocks.patchify.patchify_hiera.rst", "cyto_dl.nn.vits.decoder.rst", "cyto_dl.nn.vits.encoder.rst", "cyto_dl.nn.vits.mae.rst", "cyto_dl.nn.vits.predictor.rst", "cyto_dl.nn.vits.seg.rst", "cyto_dl.nn.vits.utils.rst", "cyto_dl.point_cloud.rst", "cyto_dl.point_cloud.io.rst", "cyto_dl.point_cloud.io.read_pcloud.rst", "cyto_dl.train.rst", "cyto_dl.utils.rst", "cyto_dl.utils.array.rst", "cyto_dl.utils.checkpoint.rst", "cyto_dl.utils.config.rst", "cyto_dl.utils.download_test_data.rst", "cyto_dl.utils.dummy_dataset.rst", "cyto_dl.utils.pylogger.rst", "cyto_dl.utils.rich_utils.rst", "cyto_dl.utils.rotation.rst", "cyto_dl.utils.spharm.rst", "cyto_dl.utils.spharm.reconstruction.rst", "cyto_dl.utils.spharm.rotation.rst", "cyto_dl.utils.template_utils.rst", "index.rst", "installation.rst", "modules.rst", "overview.md", "using_examples.rst"], "titles": ["Using the API", "cyto_dl package", "cyto_dl.api package", "cyto_dl.api.cyto_dl_model package", "cyto_dl.api.cyto_dl_model.cyto_dl_base_model module", "cyto_dl.api.cyto_dl_model.segmentation_plugin_model module", "cyto_dl.api.data module", "cyto_dl.api.model module", "cyto_dl.callbacks package", "cyto_dl.callbacks.callback_utils module", "cyto_dl.callbacks.csv_saver module", "cyto_dl.callbacks.image_saver module", "cyto_dl.callbacks.latent_walk module", "cyto_dl.callbacks.layer_freeze module", "cyto_dl.callbacks.model_utils module", "cyto_dl.callbacks.outlier_detection module", "cyto_dl.compile module", "cyto_dl.dataframe package", "cyto_dl.dataframe.readers module", "cyto_dl.dataframe.transforms package", "cyto_dl.dataframe.transforms.filter module", "cyto_dl.dataframe.transforms.group_cols module", "cyto_dl.dataframe.transforms.misc module", "cyto_dl.dataframe.transforms.split module", "cyto_dl.datamodules package", "cyto_dl.datamodules.array module", "cyto_dl.datamodules.data_dict module", "cyto_dl.datamodules.dataframe package", "cyto_dl.datamodules.dataframe.dataframe_datamodule module", "cyto_dl.datamodules.dataframe.grouped_dataframe_datamodule module", "cyto_dl.datamodules.dataframe.utils module", "cyto_dl.datamodules.folder module", "cyto_dl.datamodules.multidim_image module", "cyto_dl.datamodules.smartcache module", "cyto_dl.datamodules.torchvision module", "cyto_dl.eval module", "cyto_dl.image package", "cyto_dl.image.io package", "cyto_dl.image.io.aicsimage_loader module", "cyto_dl.image.io.monai_bio_reader module", "cyto_dl.image.io.numpy_reader module", "cyto_dl.image.io.ome_zarr_reader module", "cyto_dl.image.io.polygon_loader module", "cyto_dl.image.io.skimage_reader module", "cyto_dl.image.transforms package", "cyto_dl.image.transforms.add_meta module", "cyto_dl.image.transforms.bright_sampler module", "cyto_dl.image.transforms.clip module", "cyto_dl.image.transforms.contrastadjust module", "cyto_dl.image.transforms.generate_jepa_masks module", "cyto_dl.image.transforms.merge module", "cyto_dl.image.transforms.multiscale_cropper module", "cyto_dl.image.transforms.pad module", "cyto_dl.image.transforms.project module", "cyto_dl.image.transforms.rotation_mask_transform module", "cyto_dl.image.transforms.save module", "cyto_dl.image.transforms.so2_random_rotation module", "cyto_dl.image.transforms.track_transforms module", "cyto_dl.loggers package", "cyto_dl.loggers.mlflow module", "cyto_dl.models package", "cyto_dl.models.base_model module", "cyto_dl.models.basic_model module", "cyto_dl.models.classification package", "cyto_dl.models.classification.classification module", "cyto_dl.models.classification.timepoint_classification module", "cyto_dl.models.contrastive package", "cyto_dl.models.contrastive.contrastive module", "cyto_dl.models.handlers package", "cyto_dl.models.handlers.base_handler module", "cyto_dl.models.handlers.image_handler module", "cyto_dl.models.handlers.load_image_patch module", "cyto_dl.models.im2im package", "cyto_dl.models.im2im.gan module", "cyto_dl.models.im2im.multi_task module", "cyto_dl.models.im2im.utils package", "cyto_dl.models.im2im.utils.instance_seg module", "cyto_dl.models.im2im.utils.noise_annealer module", "cyto_dl.models.im2im.utils.postprocessing package", "cyto_dl.models.im2im.utils.postprocessing.act_thresh_label module", "cyto_dl.models.im2im.utils.postprocessing.auto_thresh module", "cyto_dl.models.im2im.utils.postprocessing.dict_to_im module", "cyto_dl.models.jepa package", "cyto_dl.models.jepa.ijepa module", "cyto_dl.models.jepa.iwm module", "cyto_dl.models.jepa.jepa_base module", "cyto_dl.models.utils package", "cyto_dl.models.utils.mlflow module", "cyto_dl.models.utils.utils module", "cyto_dl.models.vae package", "cyto_dl.models.vae.base_vae module", "cyto_dl.models.vae.image_canon_vae module", "cyto_dl.models.vae.image_encoder module", "cyto_dl.models.vae.image_vae module", "cyto_dl.models.vae.latent_loss_vae module", "cyto_dl.models.vae.o2_spharm_vae package", "cyto_dl.models.vae.o2_spharm_vae.o2_spharm_encoder module", "cyto_dl.models.vae.o2_spharm_vae.o2_spharm_vae module", "cyto_dl.models.vae.point_cloud_vae module", "cyto_dl.models.vae.priors package", "cyto_dl.models.vae.priors.abstract_prior module", "cyto_dl.models.vae.priors.gaussian module", "cyto_dl.models.vae.priors.identity_prior module", "cyto_dl.models.vae.priors.joint_prior module", "cyto_dl.models.vae.tabular_vae module", "cyto_dl.nn package", "cyto_dl.nn.discriminators package", "cyto_dl.nn.discriminators.multi_scale_discriminator module", "cyto_dl.nn.discriminators.n_layer_discriminator module", "cyto_dl.nn.head package", "cyto_dl.nn.head.base_head module", "cyto_dl.nn.head.gan_head module", "cyto_dl.nn.head.gan_head_superres module", "cyto_dl.nn.head.mae_head module", "cyto_dl.nn.head.mask_head module", "cyto_dl.nn.head.res_blocks_head module", "cyto_dl.nn.head.vic_reg module", "cyto_dl.nn.hr_skip module", "cyto_dl.nn.losses package", "cyto_dl.nn.losses.adversarial_loss module", "cyto_dl.nn.losses.chamfer_loss module", "cyto_dl.nn.losses.continuous_bernoulli module", "cyto_dl.nn.losses.cosine_loss module", "cyto_dl.nn.losses.gan_loss module", "cyto_dl.nn.losses.gaussian_nll_loss module", "cyto_dl.nn.losses.geomloss module", "cyto_dl.nn.losses.loss_wrapper module", "cyto_dl.nn.losses.spharm_loss module", "cyto_dl.nn.losses.threshold_loss module", "cyto_dl.nn.losses.vic_reg module", "cyto_dl.nn.losses.weibull module", "cyto_dl.nn.losses.weighted_mse_loss module", "cyto_dl.nn.mlp module", "cyto_dl.nn.point_cloud package", "cyto_dl.nn.point_cloud.dgcnn module", "cyto_dl.nn.point_cloud.folding_net module", "cyto_dl.nn.point_cloud.graph_functions module", "cyto_dl.nn.point_cloud.vnn module", "cyto_dl.nn.res_unit module", "cyto_dl.nn.spatial_transformer module", "cyto_dl.nn.track_sequence_predictor module", "cyto_dl.nn.vits package", "cyto_dl.nn.vits.blocks package", "cyto_dl.nn.vits.blocks.cross_attention module", "cyto_dl.nn.vits.blocks.intermediate_weigher module", "cyto_dl.nn.vits.blocks.masked_unit_attention module", "cyto_dl.nn.vits.blocks.patchify package", "cyto_dl.nn.vits.blocks.patchify.patchify module", "cyto_dl.nn.vits.blocks.patchify.patchify_base module", "cyto_dl.nn.vits.blocks.patchify.patchify_hiera module", "cyto_dl.nn.vits.decoder module", "cyto_dl.nn.vits.encoder module", "cyto_dl.nn.vits.mae module", "cyto_dl.nn.vits.predictor module", "cyto_dl.nn.vits.seg module", "cyto_dl.nn.vits.utils module", "cyto_dl.point_cloud package", "cyto_dl.point_cloud.io package", "cyto_dl.point_cloud.io.read_pcloud module", "cyto_dl.train module", "cyto_dl.utils package", "cyto_dl.utils.array module", "cyto_dl.utils.checkpoint module", "cyto_dl.utils.config module", "cyto_dl.utils.download_test_data module", "cyto_dl.utils.dummy_dataset module", "cyto_dl.utils.pylogger module", "cyto_dl.utils.rich_utils module", "cyto_dl.utils.rotation module", "cyto_dl.utils.spharm package", "cyto_dl.utils.spharm.reconstruction module", "cyto_dl.utils.spharm.rotation module", "cyto_dl.utils.template_utils module", "Overview", "Installation", "cyto_dl", "Overview", "Modifying the Example Configs"], "terms": {"cytodl": [0, 5, 173, 176], "can": [0, 18, 20, 21, 28, 29, 31, 38, 40, 51, 52, 77, 90, 153, 155, 173, 176, 177], "import": [0, 138, 173, 176, 177], "python": [0, 21, 166, 172, 174], "packag": [0, 175], "programmat": 0, "cyto_dl": [0, 173, 176, 177], "cytodlmodel": [0, 1, 2, 7], "class": [0, 4, 5, 6, 7, 10, 11, 13, 14, 15, 21, 22, 23, 25, 26, 28, 29, 30, 31, 32, 33, 34, 38, 39, 40, 41, 42, 43, 45, 46, 47, 48, 49, 50, 51, 52, 53, 55, 57, 59, 61, 62, 64, 65, 67, 73, 74, 76, 77, 79, 80, 81, 83, 84, 85, 90, 100, 101, 102, 103, 107, 108, 110, 111, 112, 113, 114, 115, 116, 117, 119, 120, 121, 122, 123, 124, 126, 128, 129, 130, 131, 132, 138, 139, 140, 143, 144, 145, 147, 148, 149, 150, 151, 152, 153, 154, 165, 168], "addit": [0, 25, 26, 28, 29, 31, 33, 39, 61, 64, 67, 73, 74, 83, 84, 85, 90, 173, 176, 177], "command": [0, 166, 167, 172], "line": [0, 166, 167, 172, 177], "interfac": [0, 90], "we": [0, 173, 176, 177], "provid": [0, 7, 22, 32, 47, 48, 76, 126, 151, 167, 173, 176, 177], "configur": [0, 4, 5, 7, 167, 177], "file": [0, 4, 7, 18, 28, 29, 31, 32, 33, 40, 110, 172, 173, 176, 177], "model": [0, 1, 2, 4, 51, 115, 140, 161, 162, 167, 172, 175, 177], "make": [0, 14, 77, 172, 177], "from": [0, 4, 5, 7, 14, 18, 25, 26, 28, 29, 30, 31, 32, 38, 39, 43, 45, 46, 49, 51, 57, 62, 77, 81, 112, 115, 120, 123, 138, 143, 153, 161, 167, 168, 172, 177], "dictionari": [0, 7, 21, 25, 26, 28, 29, 30, 31, 32, 38, 40, 45, 47, 48, 51, 61, 67, 76, 81, 90, 151, 152, 161, 177], "cfg": [0, 4, 5, 35, 159, 163, 167, 172], "data": [0, 1, 2, 7, 22, 23, 25, 26, 28, 29, 30, 31, 33, 38, 39, 41, 43, 45, 79, 83, 84, 85, 115, 161, 167, 177], "load_config_from_dict": [0, 2, 7], "yaml": [0, 4], "load_config_from_yaml": 0, "path": [0, 4, 7, 11, 18, 26, 28, 29, 31, 32, 33, 38, 40, 42, 55, 62, 87, 145, 154, 167, 172, 173, 176, 177], "name": [0, 4, 13, 18, 21, 28, 29, 30, 31, 46, 47, 48, 50, 51, 52, 64, 153, 166, 173, 176, 177], "one": [0, 22, 28, 29, 30, 61, 172, 177], "our": [0, 177], "default": [0, 4, 5, 23, 28, 29, 41, 74, 79, 83, 84, 85, 138, 145, 154, 173, 176, 177], "avail": [0, 29, 30], "option": [0, 13, 18, 20, 21, 23, 26, 28, 29, 31, 32, 42, 49, 61, 62, 67, 74, 76, 83, 84, 85, 90, 115, 145, 154, 161, 167, 172, 174, 177], "ar": [0, 18, 20, 21, 22, 23, 25, 26, 28, 29, 31, 32, 38, 39, 42, 47, 48, 49, 61, 64, 73, 74, 76, 108, 115, 123, 148, 150, 151, 152, 153, 167, 172, 173, 176, 177], "gan": [0, 2, 6, 60, 72, 111, 112, 114, 123, 173, 176, 177], "instance_seg": [0, 2, 6, 72, 75], "labelfre": [0, 6, 173, 176, 177], "segmen": 0, "when": [0, 25, 26, 28, 29, 30, 31, 61, 76, 138, 173, 176], "experi": [0, 5, 177], "overrid": [0, 7, 173, 176], "list": [0, 7, 11, 13, 18, 20, 22, 25, 28, 29, 31, 32, 38, 40, 42, 43, 45, 50, 51, 53, 57, 62, 67, 74, 76, 79, 81, 90, 116, 145, 147, 148, 149, 150, 151, 152, 153, 154, 155, 163, 172], "string": [0, 13, 18, 20, 21, 31, 61, 138], "format": [0, 38, 57, 177], "same": [0, 21, 57, 62, 90, 115, 123, 148, 150, 151, 152, 153, 155, 177], "wai": [0, 173, 176], "cli": 0, "load_default_experi": [0, 2, 7], "batch_siz": [0, 25, 26, 28, 29, 30, 31, 84, 165, 173, 176, 177], "16": [0, 6, 152, 154], "onc": [0, 177], "i": [0, 4, 18, 20, 21, 22, 23, 25, 26, 28, 29, 30, 31, 32, 33, 42, 49, 51, 52, 53, 61, 62, 76, 79, 81, 83, 84, 85, 90, 115, 123, 126, 138, 144, 150, 151, 152, 153, 154, 155, 163, 165, 168, 172, 173, 176, 177], "attribut": 0, "overridden": [0, 177], "override_config": [0, 2, 7], "which": [0, 7, 20, 21, 25, 26, 28, 29, 31, 32, 40, 61, 77, 107, 168, 172, 173, 176, 177], "take": [0, 18, 51, 53], "valu": [0, 6, 20, 21, 22, 23, 25, 28, 32, 38, 47, 48, 49, 79, 81, 83, 84, 85, 128, 138, 143, 145, 163, 165, 172, 173, 176, 177], "made": [0, 28, 61, 177], "async": 0, "pass": [0, 25, 26, 28, 29, 31, 32, 33, 38, 39, 51, 61, 64, 67, 73, 74, 77, 83, 84, 85, 90, 107, 173, 176, 177], "run_async": [0, 7], "true": [0, 7, 23, 28, 30, 32, 38, 39, 42, 45, 47, 51, 59, 76, 87, 90, 110, 111, 113, 114, 116, 128, 138, 145, 147, 148, 150, 151, 152, 153, 154, 155, 173, 176], "await": 0, "The": [0, 20, 23, 28, 29, 32, 49, 61, 62, 83, 84, 85, 123, 138, 168, 173, 176, 177], "To": [0, 173, 176], "examin": [0, 46], "print_config": [0, 2, 7], "download": 0, "exampl": [0, 32, 51, 61, 77, 115, 138, 151, 173, 176], "download_example_data": [0, 2, 7], "thi": [0, 4, 14, 23, 28, 29, 30, 32, 49, 51, 61, 62, 77, 165, 173, 176, 177], "api": [1, 175], "cyto_dl_model": [1, 2], "modul": [1, 2, 3, 8, 17, 19, 24, 27, 36, 37, 44, 58, 60, 63, 66, 68, 72, 75, 78, 82, 86, 89, 95, 99, 105, 106, 109, 118, 133, 141, 142, 146, 156, 157, 160, 169, 175], "autothresholdmethod": [1, 2, 6], "experimenttyp": [1, 2, 6], "hardwaretyp": [1, 2, 5, 6], "patchsiz": [1, 2, 5, 6], "callback": [1, 167, 172, 173, 175, 176], "callback_util": [1, 8], "csv_saver": [1, 8], "csvsaver": [1, 8, 10], "jepasav": [1, 8, 10], "image_sav": [1, 8], "imagesav": [1, 8, 11], "latent_walk": [1, 8], "layer_freez": [1, 8], "layerfreez": [1, 8, 13], "model_util": [1, 8], "getembed": [1, 8, 14], "get_all_embed": [1, 8, 14], "save_predictions_classifi": [1, 8, 14], "outlier_detect": [1, 8], "outlierdetect": [1, 8, 15], "datafram": [1, 24, 33, 173, 175, 176], "transform": [1, 17, 25, 26, 28, 29, 30, 31, 32, 33, 36, 38, 42, 43, 76, 79, 115, 140, 150, 151, 152, 153, 173, 176, 177], "reader": [1, 17], "filter_column": [1, 17, 18, 19, 20], "read_csv": [1, 17, 18], "read_datafram": [1, 17, 18], "read_h5ad": [1, 17, 18], "read_parquet": [1, 17, 18], "datamodul": [1, 14, 173, 175, 176], "arrai": [1, 24, 32, 42, 81, 115, 126, 148, 160, 173, 176], "make_array_dataload": [1, 24, 25], "data_dict": [1, 24], "make_data_dict_dataload": [1, 24, 26], "folder": [1, 24, 167, 173, 176], "make_folder_dataload": [1, 24, 31], "multidim_imag": [1, 24], "multidimimagedataset": [1, 24, 32], "make_multidim_image_dataload": [1, 24, 32], "smartcach": [1, 24], "smartcachedatamodul": [1, 24, 33], "torchvis": [1, 24], "cifar10": [1, 24, 34], "cifar100": [1, 24, 34], "clevrclassif": [1, 24, 34], "caltech101": [1, 24, 34], "caltech256": [1, 24, 34], "celeba": [1, 24, 34], "cityscap": [1, 24, 34], "cococapt": [1, 24, 34], "cocodetect": [1, 24, 34], "country211": [1, 24, 34], "dtd": [1, 24, 34], "datasetfold": [1, 24, 34], "emnist": [1, 24, 34], "eurosat": [1, 24, 34], "fer2013": [1, 24, 34], "fgvcaircraft": [1, 24, 34], "fakedata": [1, 24, 34], "fashionmnist": [1, 24, 34], "flickr30k": [1, 24, 34], "flickr8k": [1, 24, 34], "flowers102": [1, 24, 34], "food101": [1, 24, 34], "gtsrb": [1, 24, 34], "hmdb51": [1, 24, 34], "inaturalist": [1, 24, 34], "imagefold": [1, 24, 34], "imagenet": [1, 24, 34], "kmnist": [1, 24, 34], "kinet": [1, 24, 34], "kitti": [1, 24, 34], "lfwpair": [1, 24, 34], "lfwpeopl": [1, 24, 34], "lsun": [1, 24, 34], "lsunclass": [1, 24, 34], "mnist": [1, 24, 34], "movingmnist": [1, 24, 34], "omniglot": [1, 24, 34], "oxfordiiitpet": [1, 24, 34], "pcam": [1, 24, 34], "phototour": [1, 24, 34], "places365": [1, 24, 34], "qmnist": [1, 24, 34], "renderedsst2": [1, 24, 34], "sbdataset": [1, 24, 34], "sbu": [1, 24, 34], "semeion": [1, 24, 34], "stl10": [1, 24, 34], "sun397": [1, 24, 34], "svhn": [1, 24, 34], "stanfordcar": [1, 24, 34], "ucf101": [1, 24, 34], "usp": [1, 24, 34], "vocdetect": [1, 24, 34], "vocsegment": [1, 24, 34], "visiondataset": [1, 24, 34], "widerfac": [1, 24, 34], "imag": [1, 11, 26, 29, 30, 32, 33, 64, 67, 73, 74, 76, 79, 81, 83, 84, 108, 110, 111, 112, 113, 114, 115, 116, 123, 126, 140, 147, 148, 149, 150, 151, 152, 153, 173, 175, 176, 177], "io": [1, 18, 36, 73, 74, 115, 150, 156, 173, 176], "logger": [1, 166, 167, 172, 173, 175, 176], "mlflow": [1, 58, 60, 86, 173, 176], "mlflowlogg": [1, 58, 59], "classif": [1, 60], "contrast": [1, 48, 60], "handler": [1, 60], "im2im": [1, 60], "jepa": [1, 49, 60, 153], "util": [1, 24, 25, 26, 27, 28, 29, 31, 60, 72, 105, 141, 173, 175, 176, 177], "vae": [1, 60], "base_model": [1, 60], "basemodel": [1, 60, 61, 62, 64, 67, 73, 74, 83, 84, 85, 90], "basemodelmeta": [1, 60, 61], "basic_model": [1, 60], "basicmodel": [1, 60, 62], "nn": [1, 62, 64, 67, 73, 74, 79, 83, 84, 85, 90, 175], "discrimin": [1, 73, 77, 105, 111, 119, 123, 173, 176], "head": [1, 64, 67, 73, 74, 105, 145, 150, 151, 152, 153, 154, 161, 173, 176, 177], "loss": [1, 61, 62, 76, 90, 102, 105, 110, 111, 112, 113, 114, 115, 116, 173, 176, 177], "point_cloud": [1, 105, 175], "vit": [1, 105], "hr_skip": [1, 105], "hrskip": [1, 105, 117], "mlp": [1, 105, 116, 141, 142, 143, 145], "res_unit": [1, 105], "residualunit": [1, 105, 138], "spatial_transform": [1, 105], "convpoolrelu": [1, 105, 139], "stn": [1, 105, 139], "track_sequence_predictor": [1, 105], "trackclassifi": [1, 105, 140], "positionalencoding1d": [1, 105, 140], "spharm": [1, 160], "create_dataload": [1, 160, 161], "extract_array_predict": [1, 160, 161], "checkpoint": [1, 4, 15, 87, 154, 160], "load_checkpoint": [1, 160, 162], "config": [1, 4, 7, 160, 161, 167, 172, 173, 176], "is_config": [1, 160, 163], "kv_to_dict": [1, 160, 163], "remove_aux_kei": [1, 160, 163], "download_test_data": [1, 160], "delete_test_data": [1, 160, 164], "parse_s3_path": [1, 160, 164], "setup_path": [1, 160, 164], "dummy_dataset": [1, 160], "dummydatamodul": [1, 160, 165], "dummydataset": [1, 160, 165], "pylogg": [1, 160], "get_pylogg": [1, 160, 166], "rich_util": [1, 160], "enforce_tag": [1, 160, 167], "print_config_tre": [1, 160, 167], "rotat": [1, 160, 169, 177], "rotationmodul": [1, 160, 168], "get_rotation_matrix_so2": [1, 160, 168], "get_rotation_matrix_so3": [1, 160, 168], "template_util": [1, 160], "close_logg": [1, 160, 172], "extra": [1, 148, 149, 151, 152, 160, 167, 172, 174], "get_metric_valu": [1, 160, 172], "instantiate_callback": [1, 160, 172], "instantiate_logg": [1, 160, 172], "log_hyperparamet": [1, 160, 172], "save_fil": [1, 160, 172], "task_wrapp": [1, 160, 172], "compil": [1, 64, 73, 74, 173, 175, 176], "eval": [1, 175], "main": [1, 35, 120, 159, 175], "train": [1, 2, 3, 4, 5, 6, 7, 11, 23, 25, 26, 28, 29, 30, 31, 32, 49, 51, 59, 61, 62, 64, 73, 74, 76, 83, 84, 85, 87, 114, 154, 173, 175, 176], "cyto_dl_base_model": [2, 3], "cytodlbasemodel": [2, 3, 4, 5], "segmentation_plugin_model": [2, 3], "segmentationpluginmodel": [2, 3, 5], "isodata": [2, 6], "li": [2, 6], "local": [2, 6], "mean": [2, 6, 46, 76, 101, 119, 121, 122], "minimum": [2, 6, 33], "multiotsu": [2, 6], "niblack": [2, 6], "otsu": [2, 6], "sauvola": [2, 6], "triangl": [2, 6], "yen": [2, 6], "label_fre": [2, 6], "segment": [2, 6, 52, 76, 77, 154, 173, 176, 177], "segmentation_arrai": [2, 6], "segmentation_plugin": [2, 5, 6], "cpu": [2, 6, 61, 81], "gpu": [2, 6, 166, 172, 177], "larg": [2, 6, 33, 177], "medium": [2, 6], "small": [2, 6, 76], "load_config_from_fil": [2, 7], "predict": [2, 3, 4, 5, 7, 10, 25, 28, 29, 32, 61, 62, 64, 67, 83, 84, 85, 110, 111, 112, 113, 114, 115, 123, 126, 153, 161, 173, 176], "save_config": [2, 3, 4, 7], "from_default_config": [3, 4], "from_existing_config": [3, 4], "get_config": [3, 4, 60, 86, 87], "get_experiment_nam": [3, 4], "set_experiment_nam": [3, 4], "get_hardware_typ": [3, 5], "get_input_channel": [3, 5], "get_manifest_column_nam": [3, 5], "get_patch_s": [3, 5], "get_raw_image_channel": [3, 5], "get_split_column": [3, 5], "remove_split_column": [3, 5], "set_hardware_typ": [3, 5], "set_input_channel": [3, 5], "set_manifest_column_nam": [3, 5], "set_patch_s": [3, 5], "set_raw_image_channel": [3, 5], "set_split_column": [3, 5], "dictconfig": [4, 5, 26, 35, 159, 163, 167, 172], "sourc": [4, 5, 6, 7, 10, 11, 13, 14, 15, 18, 20, 21, 22, 23, 25, 26, 28, 29, 30, 31, 32, 33, 35, 38, 39, 40, 41, 42, 43, 45, 46, 47, 48, 49, 50, 51, 52, 53, 55, 57, 59, 61, 62, 64, 65, 67, 73, 74, 76, 77, 79, 80, 81, 83, 84, 85, 87, 88, 90, 100, 101, 102, 103, 107, 108, 110, 111, 112, 113, 114, 115, 116, 117, 119, 120, 121, 122, 123, 124, 126, 128, 129, 130, 131, 132, 138, 139, 140, 143, 144, 145, 147, 148, 149, 150, 151, 152, 153, 154, 155, 159, 161, 162, 163, 164, 165, 166, 167, 168, 172], "base": [4, 5, 6, 7, 10, 11, 13, 14, 15, 21, 22, 26, 28, 29, 30, 31, 32, 33, 34, 38, 39, 40, 41, 42, 43, 45, 46, 47, 48, 49, 50, 51, 52, 53, 55, 57, 59, 61, 62, 64, 65, 67, 73, 74, 76, 77, 79, 80, 81, 83, 84, 85, 90, 100, 101, 102, 103, 107, 108, 110, 111, 112, 113, 114, 115, 116, 117, 119, 120, 121, 122, 123, 124, 126, 128, 129, 130, 131, 132, 138, 139, 140, 143, 144, 145, 147, 148, 149, 150, 151, 152, 153, 154, 165, 168, 173, 176, 177], "abc": [4, 110, 148, 152], "A": [4, 5, 18, 20, 22, 26, 28, 29, 31, 62, 123, 138, 168], "us": [4, 5, 14, 18, 20, 21, 22, 23, 25, 26, 28, 29, 31, 32, 33, 38, 39, 47, 48, 49, 50, 51, 52, 61, 62, 64, 67, 73, 74, 76, 77, 83, 84, 85, 90, 110, 112, 114, 115, 123, 138, 148, 150, 151, 152, 153, 165, 167, 173, 174, 176, 177], "run": [4, 32, 74, 110, 113, 114, 116, 173, 176, 177], "cyto": [4, 174], "dl": [4, 174], "Not": [4, 114], "intend": 4, "direct": [4, 139], "client": 4, "pleas": 4, "see": [4, 18, 20, 25, 26, 28, 29, 31, 115, 138, 173, 176], "classmethod": [4, 101], "factori": 4, "method": [4, 61, 80, 115], "instead": [4, 38, 84, 152], "spatial_dim": [4, 32, 33, 49, 112, 115, 138, 145, 147, 148, 149, 150, 151, 152, 153, 154, 155, 168, 177], "int": [4, 5, 7, 11, 22, 23, 28, 29, 30, 32, 33, 46, 48, 49, 51, 52, 53, 57, 67, 76, 77, 79, 83, 84, 85, 100, 102, 107, 108, 112, 115, 116, 126, 132, 138, 145, 147, 148, 149, 150, 151, 152, 153, 154, 155, 165, 168], "return": [4, 18, 23, 32, 51, 61, 76, 107, 123, 140, 144, 163, 168], "paramet": [4, 7, 11, 13, 18, 20, 21, 22, 23, 25, 26, 28, 29, 30, 31, 32, 33, 38, 40, 42, 45, 46, 47, 48, 49, 50, 51, 52, 53, 55, 57, 62, 64, 67, 73, 74, 76, 77, 79, 81, 83, 84, 85, 90, 101, 107, 108, 110, 111, 112, 113, 114, 115, 116, 123, 126, 128, 140, 145, 148, 150, 151, 152, 153, 154, 163, 168, 172, 173, 176, 177], "dimens": [4, 32, 33, 49, 76, 115, 116, 138, 140, 144, 145, 148, 149, 150, 151, 152, 153, 154, 155], "e": [4, 18, 32, 90, 153, 173, 174, 176, 177], "g": [4, 32, 153, 173, 176, 177], "2": [4, 29, 32, 49, 51, 76, 107, 108, 117, 138, 139, 140, 145, 151, 154, 168], "config_filepath": 4, "an": [4, 6, 18, 25, 26, 29, 32, 49, 51, 53, 61, 79, 101, 126, 138, 148, 151, 173, 176, 177], "exist": [4, 23], "basi": 4, "must": [4, 13, 18, 20, 23, 25, 31, 32, 45, 51, 61, 173, 176], "gener": [4, 18, 22, 33, 49, 51, 57, 62, 76, 77, 79, 83, 84, 85, 108, 110, 111, 112, 155, 165, 173, 176], "subclass": [4, 30, 90], "want": [4, 173, 176, 177], "str": [4, 5, 7, 11, 13, 14, 15, 18, 20, 21, 22, 23, 25, 26, 28, 29, 30, 31, 32, 33, 38, 39, 40, 41, 42, 43, 45, 46, 47, 48, 50, 51, 52, 53, 55, 57, 59, 62, 64, 67, 73, 74, 76, 79, 80, 81, 83, 84, 85, 90, 114, 123, 126, 138, 147, 148, 149, 153, 154, 167, 172], "manifest_path": 4, "output_dir": [4, 7, 14], "none": [4, 5, 7, 11, 13, 14, 18, 20, 21, 22, 23, 26, 28, 29, 30, 31, 32, 33, 35, 41, 43, 51, 52, 59, 61, 62, 64, 74, 76, 79, 80, 90, 101, 108, 110, 111, 112, 113, 114, 115, 116, 117, 124, 131, 132, 138, 143, 147, 148, 149, 150, 151, 152, 153, 154, 159, 161, 165, 167, 172], "max_epoch": [4, 83, 84, 85], "handl": [5, 32, 42, 61, 123], "tupl": [5, 32, 39, 41, 43, 49, 61, 79, 138, 150, 155], "hardware_typ": 5, "input_channel": [5, 147, 148, 149, 151, 152], "target1": 5, "target2": 5, "merge_mask": 5, "exclude_mask": 5, "base_imag": 5, "patch_siz": [5, 140, 147, 148, 149, 150, 151, 152, 177], "channel": [5, 32, 33, 38, 41, 43, 46, 52, 57, 79, 108, 115, 126, 138, 148, 149, 151, 152, 154, 173, 176, 177], "split_column": [5, 28, 29, 30, 173, 176], "enum": 6, "enumer": [6, 33, 38], "threshold_isodata": 6, "threshold_li": 6, "threshold_loc": 6, "threshold_mean": 6, "threshold_minimum": 6, "threshold_multiotsu": 6, "threshold_niblack": 6, "threshold_otsu": 6, "threshold_sauvola": 6, "threshold_triangl": 6, "threshold_yen": 6, "patch": [6, 46, 49, 51, 145, 147, 148, 149, 150, 151, 152, 153, 154, 155, 177], "size": [6, 30, 49, 51, 76, 123, 138, 148, 149, 150, 151, 152, 154, 155, 177], "respect": [6, 18, 90, 177], "shape": [6, 42, 148, 165], "128": [6, 140], "64": [6, 108, 139, 149, 151, 152], "32": [6, 139, 154], "object": [7, 18, 32, 46, 76, 77, 79, 80, 81, 108, 123, 154, 168], "dict": [7, 20, 21, 25, 26, 28, 29, 32, 39, 41, 43, 51, 52, 59, 62, 73, 74, 76, 83, 84, 85, 90, 151, 152, 165, 172], "load": [7, 18, 26, 28, 29, 31, 32, 33, 38, 39, 40, 62, 173, 176, 177], "config_path": 7, "experiment_typ": 7, "directori": [7, 11, 26, 28, 29, 31, 64, 67, 73, 74, 84, 177], "float": [7, 23, 33, 46, 47, 48, 49, 76, 77, 79, 80, 83, 84, 85, 90, 123, 126, 128, 138, 145, 152, 159, 168, 172], "bool": [7, 11, 20, 23, 28, 29, 30, 32, 38, 39, 40, 41, 42, 43, 45, 47, 48, 51, 52, 53, 55, 57, 76, 79, 81, 84, 112, 115, 123, 126, 128, 138, 145, 147, 148, 150, 151, 152, 153, 154, 155, 167], "fals": [7, 11, 20, 23, 28, 29, 30, 32, 38, 39, 40, 45, 47, 48, 51, 53, 55, 57, 64, 73, 74, 76, 79, 81, 84, 90, 100, 101, 103, 108, 112, 114, 115, 117, 119, 138, 143, 145, 148, 150, 151, 152, 153, 167], "save": [7, 8, 11, 32, 36, 44, 57, 62, 64, 67, 73, 74, 84, 110, 113, 114, 116, 151, 172, 173, 176, 177], "current": [7, 29, 32, 42], "on_predict_epoch_end": [8, 10], "pred_to_csv": [8, 10], "save_feat": [8, 10], "on_predict_batch_end": [8, 11, 15], "on_test_batch_end": [8, 11, 15], "on_train_batch_end": [8, 11], "on_validation_batch_end": [8, 11], "setup": [8, 13, 24, 33, 172], "on_test_epoch_end": [8, 14], "calculate_mahalanobis_hook": [8, 15], "flatten_activ": [8, 15], "on_load_checkpoint": [8, 15], "on_predict_epoch_start": [8, 15], "on_save_checkpoint": [8, 15], "on_test_epoch_start": [8, 15], "on_train_epoch_start": [8, 15], "update_covariance_hook": [8, 15], "save_dir": [10, 11, 15, 59, 64, 67, 73, 74, 84, 177], "meta_kei": [10, 45, 67], "trainer": [10, 11, 13, 14, 15, 167, 173, 176], "pl_modul": [10, 11, 13, 14, 15], "pred": [10, 14, 61, 107, 120], "stage": [10, 11, 13, 33, 61, 62, 64, 67, 73, 74, 83, 84, 85, 90, 110, 111, 113, 114, 116], "save_every_n_epoch": 11, "1": [11, 23, 29, 30, 33, 49, 51, 53, 64, 73, 74, 76, 79, 90, 101, 110, 111, 112, 113, 115, 116, 123, 126, 129, 131, 132, 138, 145, 147, 148, 149, 151, 152, 164, 177], "val": [11, 23, 28, 29, 61, 62, 88], "save_input": 11, "after": [11, 45, 76, 79, 115, 143, 172], "postprocess": [11, 72, 75, 110, 111, 112, 113, 114, 115, 116, 173, 176, 177], "ead": 11, "union": [11, 13, 18, 20, 21, 25, 26, 28, 29, 31, 32, 33, 40, 42, 50, 53, 62, 74, 76, 81], "frequenc": [11, 64, 67, 73, 74], "whether": [11, 23, 28, 29, 30, 32, 38, 39, 40, 42, 46, 47, 48, 51, 52, 53, 57, 64, 73, 74, 76, 79, 81, 84, 114, 115, 128, 138, 148, 150, 151, 152, 165, 167, 168, 177], "input": [11, 20, 22, 23, 25, 26, 28, 29, 31, 40, 62, 64, 73, 74, 76, 83, 85, 108, 110, 111, 112, 113, 114, 115, 116, 119, 121, 122, 123, 124, 131, 138, 145, 148, 149, 151, 152, 153, 163, 167, 173, 176, 177], "output": [11, 15, 32, 38, 62, 79, 90, 110, 115, 116, 123, 138, 145, 151, 154, 161, 167, 172, 173, 176], "batch": [11, 15, 21, 26, 28, 29, 30, 31, 32, 45, 55, 61, 62, 64, 65, 67, 73, 74, 83, 84, 85, 90, 101, 110, 111, 113, 114, 116, 145, 149, 165, 168, 177], "batch_idx": [11, 15, 61, 62, 64, 65, 67, 73, 74, 83, 84, 85, 90], "dataloader_idx": [11, 15], "0": [11, 15, 33, 41, 46, 47, 48, 49, 51, 76, 77, 83, 84, 85, 90, 101, 112, 115, 117, 123, 126, 128, 138, 143, 145, 147, 148, 149, 151, 152, 154, 168], "step": [11, 32, 59, 61, 77, 177], "contain": [13, 18, 20, 21, 25, 26, 28, 31, 32, 33, 42, 177], "exclud": [13, 18, 20, 31, 76], "search": 13, "within": [13, 48, 151], "layer": [13, 108, 115, 116, 123, 138, 144, 145, 150, 151, 152, 153, 154], "freez": 13, "x_label": [14, 62, 90], "id_label": [14, 61, 90], "arg": [14, 25, 26, 28, 29, 31, 34, 38, 40, 61, 62, 64, 65, 67, 73, 74, 83, 84, 85, 90, 138, 165, 167], "id_field": 14, "lightningmodul": [14, 61, 172], "train_dataload": [14, 24, 27, 28, 33, 160, 165], "val_dataload": [14, 24, 27, 28, 33, 160, 165], "test_dataload": [14, 24, 27, 28, 33, 160, 165], "todo": [14, 23, 28, 29], "better": [14, 148, 150, 151, 152, 153, 177], "mayb": 14, "vol": 14, "predictor": [14, 83, 84, 85, 105, 141], "code": [14, 173, 176], "drop": [14, 30, 143], "unnecessari": 14, "index": [14, 53, 148, 149, 150, 155], "n_epoch": 15, "layer_nam": 15, "act": [15, 138], "set": [15, 28, 48, 172], "forward": [15, 60, 61, 62, 63, 64, 66, 67, 72, 74, 82, 85, 89, 90, 99, 100, 101, 102, 103, 105, 106, 107, 108, 109, 110, 111, 112, 115, 117, 118, 119, 120, 121, 122, 124, 129, 130, 131, 132, 138, 139, 140, 141, 142, 143, 144, 145, 146, 148, 150, 151, 152, 153, 154], "hook": 15, "callabl": [15, 25, 26, 31, 32, 51, 62, 79, 112, 115, 172], "filter": [17, 18, 19, 31, 108, 154], "filter_row": [17, 19, 20], "group_col": [17, 19], "groupcol": [17, 19, 21], "misc": [17, 19], "append_class_weight": [17, 19, 22], "append_labels_to_integ": [17, 19, 22], "append_one_hot": [17, 19, 22], "make_random_df": [17, 19, 22], "split": [17, 19, 28, 29, 30, 33, 61, 173, 176, 177], "sample_n_each": [17, 19, 23], "split_datafram": [17, 19, 23], "columns_to_filt": 18, "sequenc": [18, 20, 22, 25, 26, 28, 29, 30, 31, 39, 40, 41, 43, 45, 51, 55, 62, 76, 79, 90, 138, 140, 147, 148, 149, 155, 167, 173, 176], "regex": [18, 20, 31], "startswith": [18, 20, 31], "endswith": [18, 20, 31], "column": [18, 20, 21, 22, 23, 28, 29, 30, 31, 32, 33, 173, 176, 177], "combin": [18, 30, 173, 176, 177], "differ": [18, 20, 23, 38, 123, 126, 153, 154, 173, 176], "queri": [18, 143, 145, 151, 152], "pattern": 18, "If": [18, 20, 22, 25, 28, 29, 30, 31, 32, 38, 39, 45, 62, 115, 145, 148, 150, 151, 152, 153, 155, 173, 176, 177], "suppli": [18, 20], "preced": 18, "remain": [18, 20, 21], "argument": [18, 22, 25, 26, 28, 29, 31, 33, 39, 64, 67, 73, 74, 83, 84, 85, 90, 107, 138, 173, 176], "ignor": [18, 20, 28, 29, 42, 172], "otherwis": [18, 62, 177], "logic": [18, 61], "AND": 18, "appli": [18, 25, 26, 31, 33, 51, 52, 53, 57, 79, 115, 144, 172, 173, 176], "all": [18, 23, 28, 29, 32, 33, 51, 74, 79, 150, 153, 172, 173, 176, 177], "condit": 18, "regular": [18, 20, 31], "express": [18, 20, 31], "match": [18, 20, 31, 61, 155], "substr": [18, 20, 31], "start": [18, 20, 31, 32, 33, 51, 57, 172, 173, 176], "end": [18, 20, 31], "include_column": 18, "read": [18, 32, 33, 37, 38, 39, 41, 43], "store": [18, 26, 28, 29, 31, 62, 173, 176], "csv": [18, 29, 32, 33, 38, 67, 173, 176, 177], "includ": [18, 20, 23, 148, 149, 151, 152], "onli": [18, 20, 25, 26, 28, 29, 30, 31, 32, 42, 52, 62, 67, 76, 110, 138, 150, 152, 172, 177], "given": [18, 20, 22, 23, 28, 29, 46, 61, 102, 128, 155, 173, 176], "upath": [18, 26, 28, 29, 31], "desir": 18, "result": [18, 20, 61, 62, 177], "type": [18, 21, 29, 30, 32, 38, 61, 79, 112, 115, 123, 138, 145, 163, 168, 173, 176, 177], "pd": [18, 20, 22, 23], "required_column": 18, "parquet": [18, 28, 29], "assert": 18, "expect": [18, 26, 173, 176, 177], "requir": [18, 173, 174, 176, 177], "either": [18, 20, 21, 23, 28, 61, 155, 168], "support": [18, 29, 42, 77, 123], "aren": 18, "t": [18, 33, 38, 102, 174], "found": [18, 173, 176], "valueerror": 18, "thrown": 18, "those": [18, 173, 176], "get": [18, 33, 173, 176], "append": [18, 22, 23], "without": [18, 155], "duplic": 18, "back": 18, "anndata": 18, "h5ad": 18, "r": [18, 174], "": [18, 20, 28, 29, 30, 73, 74, 76, 77, 90, 111, 112, 168, 173, 176], "doc": [18, 25, 26, 28, 29, 31, 73, 74, 115, 173, 176], "detail": [18, 20], "http": [18, 25, 26, 28, 29, 31, 39, 43, 73, 74, 115, 120, 121, 150, 173, 174, 176], "readthedoc": 18, "en": [18, 73, 74, 115, 173, 176], "latest": 18, "html": [18, 25, 26, 28, 29, 31, 73, 74, 115, 173, 176], "select": [20, 23, 30, 46, 173, 176], "dataset": [20, 25, 28, 29, 30, 32, 33, 165, 173, 176], "more": [20, 25, 26, 28, 29, 31, 172, 173, 176], "oper": 20, "It": [20, 28, 101, 177], "panda": [20, 23], "case": [20, 21, 23, 28, 29, 61, 173, 176, 177], "explicit": [20, 130], "keep": [20, 76, 148, 173, 176], "row": [20, 22, 23], "where": [20, 21, 29, 32, 61, 79, 150, 154, 163, 165, 173, 176], "group": [21, 29, 149, 168], "dtype": [21, 38, 79, 126], "kei": [21, 26, 28, 29, 30, 31, 32, 38, 40, 42, 45, 46, 47, 48, 50, 51, 52, 53, 55, 57, 61, 62, 64, 67, 73, 74, 76, 81, 83, 84, 85, 90, 143, 145, 151, 152, 163, 165, 173, 176], "becom": [21, 165, 173, 176], "kwarg": [21, 33, 34, 38, 40, 61, 62, 64, 65, 67, 73, 74, 83, 84, 85, 90, 100, 101, 102, 103, 107, 120, 165], "col": 21, "filter_col": 21, "singl": [21, 25, 28, 32, 151, 155, 173, 176], "np": [21, 25, 79], "numpi": [21, 25, 32, 38, 42, 79, 81, 161], "recogn": 21, "turn": 21, "add": [22, 45, 61, 145], "weight": [22, 62, 76, 83, 84, 85, 90, 111, 112, 126, 131, 144, 151, 154], "modifi": [22, 23, 29, 107], "its": [22, 46, 167], "integ": [22, 32, 51], "encod": [22, 49, 83, 84, 85, 89, 90, 105, 140, 141, 143, 148, 149, 150, 152], "convert": [22, 32, 42, 76, 81, 147, 148, 149, 151, 161], "hot": 22, "call": [22, 28, 29, 61, 77, 172], "function": [22, 32, 51, 62, 76, 90, 110, 111, 112, 113, 114, 115, 116, 126, 128, 172, 173, 176, 177], "one_hot_encod": 22, "b": [22, 130], "c": [22, 38], "d": 22, "n_row": 22, "100": [22, 51, 76, 83, 84, 85, 111, 112], "random": [22, 23, 28, 29, 51], "test": [22, 23, 28, 29, 32, 61, 62, 173, 176], "wrangl": 22, "pipelin": 22, "creat": [22, 25, 26, 28, 29, 31, 32, 42, 64, 76, 123, 154, 161, 173, 174, 176, 177], "number": [22, 23, 28, 29, 33, 38, 49, 51, 52, 57, 77, 83, 84, 85, 107, 108, 115, 126, 138, 145, 148, 149, 150, 151, 152, 153, 154, 155, 172, 177], "forc": 23, "seed": [23, 28, 29], "42": [23, 28, 29], "have": [23, 45, 51, 61, 76, 138, 150, 151, 152], "equal": [23, 126], "per": [23, 28, 29, 51, 76, 126, 145, 161], "In": [23, 61, 168, 173, 176, 177], "ha": [23, 77, 123, 151, 155], "less": 23, "than": [23, 138], "correspond": [23, 29, 31, 51, 165], "sampl": [23, 25, 26, 28, 29, 30, 31, 46, 49, 51, 99, 100, 101, 103, 165], "replac": [23, 33], "uniqu": [23, 173, 176], "toggl": [23, 105, 139], "upsampl": [23, 115, 154, 173, 176], "smaller": 23, "train_frac": 23, "val_frac": 23, "return_split": 23, "perform": [23, 115, 151], "three": [23, 28], "identifi": 23, "each": [23, 25, 26, 28, 29, 30, 31, 33, 49, 57, 61, 64, 76, 90, 138, 140, 144, 148, 149, 150, 151, 152, 153, 154, 155, 165, 173, 176, 177], "belong": 23, "extend": 23, "enabl": 23, "balanc": 23, "stratifi": 23, "fraction": [23, 152], "valid": [23, 33, 61, 64, 73, 74, 173, 176], "By": [23, 177], "half": 23, "between": [23, 29, 30, 32, 64, 73, 74, 101, 110, 115, 126, 153], "separ": [23, 32, 38, 76], "reproduc": 23, "dataframe_datamodul": [24, 27], "dataframedatamodul": [24, 27, 28, 29], "grouped_dataframe_datamodul": [24, 27], "groupeddataframedatamodul": [24, 27, 29], "alternatingbatchsampl": [24, 27, 29, 30], "removenankeysd": [24, 27, 30], "get_canonical_split_nam": [24, 27, 30], "get_dataset": [24, 27, 28, 30], "make_multiple_dataframe_split": [24, 27, 30], "make_single_dataframe_split": [24, 27, 30], "parse_transform": [24, 27, 30], "create_metatensor": [24, 32], "get_per_file_arg": [24, 32, 33], "is_batch": [24, 32], "make_dataload": [24, 27, 28, 29, 33], "predict_dataload": [24, 27, 28, 33, 160, 165], "prepare_data": [24, 33], "ndarrai": [25, 39, 41, 43, 76, 81], "source_kei": [25, 84], "dataloader_kwarg": [25, 26, 28, 29, 31, 32], "dataload": [25, 26, 28, 29, 31, 32, 33, 62, 161, 173, 176], "element": [25, 61, 138, 155, 165], "keyword": [25, 26, 28, 29, 31, 39], "torch": [25, 26, 28, 29, 30, 31, 42, 62, 64, 73, 74, 79, 90, 101, 108, 111, 112, 115, 123, 143, 144, 145, 168, 177], "instanti": [25, 26, 28, 29, 31, 90, 172, 173, 176], "asid": [25, 26, 28, 29, 31], "shuffl": [25, 26, 28, 29, 30, 31], "among": [25, 26, 28, 29, 31], "num_work": [25, 26, 28, 29, 31, 33, 173, 176], "etc": [25, 26, 28, 29, 31, 173, 176, 177], "pytorch": [25, 26, 28, 29, 30, 31, 62, 90, 173, 176], "info": [25, 26, 28, 29, 31], "org": [25, 26, 28, 29, 31, 121], "stabl": [25, 26, 28, 29, 31, 73, 74, 115, 173, 176], "cache_dir": [26, 28, 29, 30, 31, 173, 176], "usual": 26, "cach": [26, 28, 29, 31, 33, 173, 176], "acceler": [26, 28, 29, 31], "get_dataload": [27, 28, 160, 165], "split_map": [28, 29, 30], "just_infer": [28, 29, 30], "subsampl": [28, 29], "refresh_subsampl": [28, 29], "smartcache_arg": [28, 29, 30], "lightningdatamodul": [28, 33, 165], "lightn": [28, 62, 172, 173, 176, 177], "fold": 28, "addition": [28, 172], "go": 28, "flag": 28, "so": [28, 29, 77, 168, 173, 176, 177], "whole": [28, 29], "predict_datamodul": 28, "simpli": [28, 61, 173, 176], "construct": [28, 173, 176], "regardless": 28, "specif": [28, 29, 61, 73, 74, 90, 115, 173, 176, 177], "everyth": [28, 29, 48], "document": [28, 29, 173, 176], "just": [28, 29, 84, 151], "infer": [28, 29, 61, 73, 74, 90, 100, 101, 103, 168], "epoch": [28, 29, 83, 84, 85], "refresh": [28, 29], "time": [28, 29, 30, 77, 151, 152, 172], "smartcachedataset": [28, 29], "target_column": [29, 30], "grouping_column": [29, 30], "should": [29, 32, 33, 40, 51, 61, 115, 126, 138, 151, 152, 154, 177], "some": [29, 32, 138, 173, 176, 177], "criterion": 29, "leverag": 29, "two": [29, 50, 101, 173, 176], "multitask": 29, "ground": [29, 30, 76, 115, 123, 173, 176, 177], "truth": [29, 30, 76, 115, 123, 173, 176, 177], "task": [29, 30, 64, 67, 73, 74, 77, 110, 111, 113, 114, 115, 116, 147, 148, 149, 172, 173, 176, 177], "characterist": 29, "altern": [29, 30, 177], "dure": [29, 32, 64, 67, 73, 74, 114, 154, 172, 173, 176], "factor": [29, 51, 154, 177], "homogen": 29, "across": [29, 53, 76, 173, 176], "subset": 30, "drop_last": 30, "sampler": 30, "subsetrandomsampl": 30, "batchsampl": 30, "mutual": 30, "exclus": 30, "monai": [30, 33, 51, 53, 73, 74, 115, 138, 173, 176, 177], "wrap": [30, 62, 172, 173, 176], "repres": [30, 138], "last": [30, 52, 87, 123, 138, 144, 151], "incomplet": 30, "randomli": [30, 49, 52, 138], "follow": [30, 151, 152, 177], "order": [30, 138, 148, 149, 152, 154, 167], "remov": [30, 49, 76, 152], "nan": 30, "ad": [30, 45, 49, 177], "allow_missing_kei": [30, 38, 47, 48, 51, 52, 53, 55, 57, 76, 81], "allow": [30, 38, 52, 55, 57, 177], "multi": [30, 32, 151, 166, 172, 173, 176, 177], "target": [30, 49, 61, 67, 77, 84, 119, 121, 122, 123, 124, 128, 130, 131, 153, 173, 176, 177], "split_path": 30, "dataframe_path": 30, "orig_fnam": 31, "origin": [31, 45], "filenam": [31, 39, 41, 43], "extens": 31, "out": [31, 64, 73, 74, 138, 152, 154, 173, 176], "rule": [31, 177], "csv_path": [32, 33], "img_path_column": [32, 33], "channel_column": [32, 33], "out_kei": [32, 38], "3": [32, 33, 49, 51, 76, 77, 107, 108, 112, 115, 117, 138, 145, 147, 148, 149, 150, 151, 152, 153, 154, 168, 174, 177], "scene_column": 32, "scene": [32, 33, 38, 39], "time_start_column": 32, "time_stop_column": 32, "stop": 32, "time_step_column": 32, "dict_meta": 32, "dask_load": [32, 38, 39], "dimension": [32, 100, 101, 102, 151], "timelaps": 32, "metadata": [32, 45, 165], "timepoint": [32, 33, 38, 39], "extract": [32, 38, 39, 83, 84, 85, 152], "spatial": [32, 33, 49, 76, 107, 115, 126, 138, 145, 148, 149, 151, 152, 153, 154, 177], "yx": [32, 49, 148, 149, 152, 177], "zyx": [32, 49, 148, 149, 152, 154, 177], "specifi": [32, 33, 38, 39, 48, 53, 81, 90, 149, 151, 152, 173, 176, 177], "multipl": [32, 51, 107, 108, 138, 153, 173, 176], "thei": [32, 45, 77], "comma": [32, 38], "scene1": 32, "scene2": 32, "ani": [32, 59, 79, 126, 173, 176, 177], "start_column": 32, "stop_column": 32, "step_column": 32, "For": [32, 51, 115, 138, 151, 173, 176, 177], "everi": [32, 62, 163], "other": [32, 173, 176, 177], "version": [32, 107, 177], "accept": [32, 126], "normal": [32, 108, 138, 143, 144, 145, 173, 176, 177], "could": 32, "here": [32, 33, 61, 121, 173, 176, 177], "dask": [32, 38, 39], "full": [32, 38, 39], "memori": [32, 33, 38, 39, 177], "befor": [32, 38, 39, 51, 57, 128, 172], "img": [32, 39, 41, 43, 57, 76, 81, 140, 148, 151, 152, 154], "meta": 32, "df": [32, 33], "x": [32, 61, 62, 64, 74, 85, 90, 107, 108, 110, 111, 112, 115, 117, 120, 129, 132, 138, 139, 143, 144, 145, 149, 151, 177], "listconfig": [32, 52, 57, 163], "multidimimag": 32, "cannot": 32, "dim": [32, 76, 107, 108, 116, 132, 145], "multiscen": 32, "alreadi": [32, 173, 176], "compos": [33, 167], "img_data": 33, "n_val": 33, "20": 33, "pct_val": 33, "raw": [33, 38, 51, 173, 176, 177], "ch": [33, 46, 79], "num_neighbor": 33, "4": [33, 140, 143, 145, 150, 151, 152, 154], "cache_r": 33, "5": [33, 49, 51, 76, 129], "replace_r": 33, "czi": 33, "don": 33, "fit": 33, "bioio": 33, "n_imag": 33, "percentag": 33, "neighbor": [33, 76], "worker": 33, "most": [33, 61], "schedul": 33, "parallel": 33, "submodul": [36, 156, 175], "aicsimage_load": [36, 37], "aicsimageloaderd": [36, 37, 38], "monai_bio_read": [36, 37, 38], "monaibioread": [36, 37, 39], "numpy_read": [36, 37], "readnumpyfil": [36, 37, 40], "ome_zarr_read": [36, 37], "omezarrread": [36, 37, 41], "polygon_load": [36, 37], "polygonloaderd": [36, 37, 42], "skimage_read": [36, 37], "skimageread": [36, 37, 43], "add_meta": [36, 44], "addmeta": [36, 44, 45], "bright_sampl": [36, 44], "brightsampl": [36, 44, 46], "clip": [36, 44, 48], "clipd": [36, 44, 47], "contrastadjust": [36, 44], "contrastadjustd": [36, 44, 48], "generate_jepa_mask": [36, 44], "jepamaskgener": [36, 44, 49], "merg": [36, 44], "multiscale_cropp": [36, 44], "randommultiscalecropd": [36, 44, 51], "pad": [36, 44, 76, 138], "padzd": [36, 44, 52], "project": [36, 44, 145, 173, 176], "maxprojectd": [36, 44, 53], "rotation_mask_transform": [36, 44], "so2_random_rot": [36, 44], "track_transform": [36, 44], "cropres": [36, 44, 57], "generatetracklabel": [36, 44, 57], "perchannel": [36, 44, 57], "splittrackd": [36, 44, 57], "split_arg": [37, 38], "get_data": [37, 39, 41, 43], "verify_suffix": [37, 39, 41, 43], "path_kei": 38, "scene_kei": 38, "kwargs_kei": 38, "dimension_order_out": [38, 177], "float16": 38, "channel_kei": 38, "timepoint_kei": 38, "fix": [38, 148, 150, 151, 152, 153, 155], "initi": [38, 77, 115, 123, 166], "bioimag": [38, 39], "get_image_dask_data": [38, 39], "miss": [38, 42, 47, 48, 52, 53, 55, 57, 76, 81], "adapt": [39, 43, 120], "github": [39, 43, 120, 150, 174], "com": [39, 43, 120, 174], "mmv": [39, 43], "lab": [39, 43], "mmv_im2im": [39, 43], "reader_kwarg": 39, "imageread": [39, 41, 43], "get_image_data": 39, "pathlik": [39, 41, 43], "maptransform": 40, "thereof": 40, "interpret": 40, "point": [40, 51, 76, 138, 177], "cloud": [40, 120], "remot": 40, "fsspec": 40, "locat": [40, 173, 176], "level": [41, 149, 173, 176, 177], "image_nam": 41, "shape_reference_kei": 42, "propagate_3d": 42, "missing_key_mod": 42, "rais": [42, 53, 76, 81], "polygon": 42, "vertic": 42, "tensor": [42, 81, 84, 85, 101, 123, 124, 131, 138, 145, 147, 154, 168, 177], "mask": [42, 49, 50, 85, 114, 145, 147, 148, 149, 150, 151, 152, 154], "propag": 42, "2d": [42, 49, 83, 115, 138, 148, 149, 168, 173, 176, 177], "3d": [42, 49, 76, 83, 115, 138, 148, 149, 168, 173, 176, 177], "how": [42, 115, 165, 173, 176], "keyerror": 42, "do": [42, 115, 123, 173, 176], "noth": 42, "new": [42, 45, 173, 176, 177], "blank": 42, "remove_excess_pixel": [44, 49], "generate_slic": [44, 51], "image_kei": [45, 50, 52, 76], "delet": 45, "updat": [45, 61], "metatensor": [45, 165], "been": [45, 151], "threshold": [46, 79, 128], "base_prob": 46, "decid": [46, 51], "intens": [46, 47, 48, 79, 177], "background": [46, 48], "rate": [46, 138, 145], "probabilii": 46, "check": [46, 52], "against": 46, "low": [47, 48, 49], "01": 47, "high": [47, 48, 49], "99": 47, "percentil": [47, 48], "absolut": [47, 48], "lower": [47, 48], "bound": [47, 48], "upper": [47, 48], "per_channel": 47, "resiz": [47, 48, 51, 57, 177], "fail": [47, 48], "adjust": 48, "rang": [48, 49, 79, 173, 176, 177], "outsid": 48, "everywher": 48, "mask_siz": 49, "12": [49, 145, 151, 152, 153, 154], "block_aspect_ratio": 49, "num_patch": [49, 148, 150, 151, 152, 153, 154, 155], "6": [49, 76, 154], "24": 49, "mask_ratio": [49, 147, 148, 149, 151, 152, 154], "9": [49, 174], "randomizabletransform": [49, 51, 57, 76], "block": [49, 105, 115, 138, 141, 151, 152], "contigu": 49, "work": [49, 101, 148, 150, 151, 152, 153, 174, 177], "until": [49, 177], "met": 49, "exceed": [49, 177], "exterior": 49, "exactli": 49, "determin": [49, 115, 138, 167, 173, 176, 177], "aspect": 49, "ratio": [49, 115, 138, 145, 152, 154], "proport": 49, "pixel": [49, 57, 76, 126, 148, 149, 151, 152, 177], "along": 49, "boundari": 49, "reach": 49, "mask_kei": [50, 114], "base_image_kei": 50, "output_nam": 50, "serv": 50, "patch_shap": [51, 139, 177], "scales_dict": 51, "patch_per_imag": 51, "selection_fn": 51, "max_attempt": 51, "style": [51, 53, 152, 177], "slice": [51, 52, 76, 177], "scale": [51, 107, 111, 112, 123, 126], "superresolut": 51, "x_kei": [51, 64, 73, 74, 83, 85], "network": [51, 62, 64, 73, 74, 90, 115, 138, 177], "Its": 51, "resolut": [51, 112, 115, 151], "len": 51, "map": [51, 110, 138, 177], "seg": [51, 76, 105, 141, 173, 176], "would": [51, 61], "kept": 51, "discard": 51, "max": [51, 53], "attempt": 51, "try": [51, 177], "quit": 51, "image_dict": 51, "segmentation_kei": 52, "pad_amount": 52, "pad_kei": 52, "segmentation_ch": 52, "top": [52, 154, 177], "bottom": 52, "crop": 52, "repeat": [52, 151, 152, 155], "first": [52, 101, 108, 151, 168, 173, 176, 177], "present": [52, 76], "offset": 52, "presenc": 52, "projection_dim": 53, "nczyx": 53, "comput": [53, 101, 128, 168], "error": [53, 76, 81], "save_path": 55, "disk": 55, "debug": 55, "max_shift": 57, "8": [57, 101, 140, 143, 145, 149, 151, 152, 154, 173, 176], "maximum": [57, 83, 84, 85], "shift": 57, "img_kei": 57, "formation_kei": 57, "breakdown_kei": 57, "breakdown": 57, "track_start_kei": 57, "track_start": 57, "label_kei": [57, 76], "label": [57, 64, 67, 76, 79, 123], "track": 57, "after_save_checkpoint": [58, 59], "log_hyperparam": [58, 59], "log_metr": [58, 59], "experiment_nam": [59, 173, 176], "lightning_log": 59, "run_nam": [59, 173, 176], "tracking_uri": [59, 87], "tag": [59, 153, 167, 172], "mlrun": 59, "prefix": 59, "artifact_loc": 59, "run_id": [59, 87], "fault_toler": 59, "ckpt_callback": 59, "param": [59, 110], "namespac": 59, "mode": [59, 87, 100, 101, 102, 103, 117, 121, 130, 172], "metric": [59, 61, 110, 113, 114, 116, 172], "timepoint_classif": [60, 63], "timepointclassif": [60, 63, 65], "base_handl": [60, 68], "image_handl": [60, 68], "load_image_patch": [60, 68], "multi_task": [60, 72], "multitaskim2im": [60, 72, 73, 74], "ijepa": [60, 82, 85], "iwm": [60, 82, 85], "jepa_bas": [60, 82], "jepabas": [60, 82, 83, 84, 85], "load_model_from_checkpoint": [60, 86, 87], "find_indic": [60, 86, 88], "o2_spharm_va": [60, 89], "prior": [60, 89, 90], "base_va": [60, 89], "baseva": [60, 89, 90], "image_canon_va": [60, 89], "image_encod": [60, 89], "image_va": [60, 89], "latent_loss_va": [60, 89], "point_cloud_va": [60, 89], "tabular_va": [60, 89], "compute_metr": [60, 61], "configure_optim": [60, 61, 72, 73, 74, 82, 85], "model_step": [60, 61, 62, 63, 64, 66, 67, 72, 73, 74, 82, 83, 84, 85, 89, 90], "on_train_start": [60, 61], "parse_batch": [60, 61], "predict_step": [60, 61, 63, 64, 65, 66, 67, 72, 73, 74, 82, 84, 85], "test_step": [60, 61, 72, 74, 82, 84], "training_step": [60, 61, 72, 74], "validation_step": [60, 61, 72, 74], "log": [61, 62, 101, 172, 173, 176, 177], "assumpt": 61, "_step": 61, "whose": 61, "mai": 61, "self": [61, 76, 83, 84, 151, 152], "structur": [61, 167, 177], "part": [61, 90, 172, 173, 176], "arbitrari": 61, "denot": 61, "you": [61, 173, 176, 177], "implement": [61, 173, 176, 177], "process": [61, 172, 177], "variabl": [61, 173, 176], "hparam": 61, "id": 61, "detach": [61, 75, 78, 81, 110, 111, 112, 113, 114, 115], "consist": [61, 154, 173, 176], "your": [61, 173, 176, 177], "might": 61, "wish": 61, "post": 61, "minim": 62, "wrapper": [62, 126, 128], "around": [62, 148, 149, 151, 152], "assum": [62, 168, 173, 176, 177], "gt": [62, 120], "optim": [62, 111, 112, 173, 176, 177], "retriev": [62, 172], "adam": 62, "save_predict": 62, "serotini": [62, 173, 176], "fields_to_log": 62, "field": [62, 167, 177], "pretrained_weight": 62, "pretrain": [62, 154], "via": [62, 151], "load_state_dict": 62, "save_imag": [63, 64], "should_save_imag": [63, 64], "share": [64, 73, 74, 177], "save_images_every_n_epoch": [64, 73, 74], "base_kwarg": [64, 67, 73, 74, 83, 84, 85, 90], "logit": [64, 115, 173, 176], "text": 64, "overlaid": 64, "plot_class": [66, 67], "plot_neighbor": [66, 67], "backbon": [67, 73, 74, 110, 113, 114, 115, 116, 154, 173, 176, 177], "task_head": [67, 73, 74, 161, 177], "anchor_kei": 67, "anchor": 67, "positive_kei": 67, "posit": [67, 140, 147, 148, 149, 150, 151, 152, 153, 155], "target_kei": [67, 84, 90], "visual": [67, 173, 176], "viz_freq": 67, "x1": 67, "x2": 67, "embedding1": 67, "embedding2": 67, "instancesegclust": [72, 75, 76], "instancesegloss": [72, 75, 76], "instancesegpreprocessd": [72, 75, 76], "instancesegrandflipd": [72, 75, 76], "pad_slic": [72, 75, 76], "noise_anneal": [72, 75, 108], "noiseanneal": [72, 75, 77], "get_n_postprocess_imag": [72, 74], "get_per_head": [72, 74], "run_forward": [72, 74, 110, 111, 113, 114, 116], "basic": [73, 90, 173, 176, 177], "inference_arg": [73, 74], "slide": [73, 74], "window": [73, 74], "inference_head": 74, "run_head": [74, 109, 110, 111, 112, 113, 114, 116], "n_postprocess": [74, 110, 111, 113, 114, 116], "act_thresh_label": [75, 78], "actthreshlabel": [75, 78, 79], "auto_thresh": [75, 78], "autothreshold": [75, 78, 80], "dict_to_im": [75, 78], "dicttoim": [75, 78, 81], "cluster_object": [75, 76], "kd_cluster": [75, 76], "remove_small_skeleton": [75, 76], "embed_from_skel": [75, 76], "keep_largest_cc": [75, 76], "label_2d": [75, 76], "shrink": [75, 76], "skeleton_tal": [75, 76], "smooth_embed": [75, 76], "topology_preserving_thin": [75, 76], "update_nois": [75, 77], "anisotropi": 76, "skel_threshold": 76, "semantic_threshold": 76, "min_siz": 76, "1000": 76, "distance_threshold": 76, "progress": 76, "cluster": 76, "instanceseg": 76, "find": 76, "skeleton": 76, "assign": 76, "semant": 76, "embed": [76, 84, 145, 147, 148, 149, 150, 151, 152, 153, 154, 155], "nearest": [76, 117], "distanc": [76, 177], "skel": 76, "closest": 76, "below": [76, 128, 173, 176], "touch": 76, "edg": [76, 177], "compon": [76, 167], "kernel_s": [76, 138, 139], "thin": 76, "keep_largest": 76, "instanc": [76, 77, 138, 154, 177], "kernel": [76, 138], "gaussian": [76, 77, 89, 99], "smooth": 76, "flow": 76, "amount": [76, 115, 138], "psuedo": 76, "largest": 76, "connect": [76, 115, 138, 154], "iseg": 76, "vector": [76, 168, 177], "z": [76, 90, 100, 101, 102, 152, 168, 173, 176], "im": [76, 108], "topologi": 76, "preserv": 76, "binari": 76, "max_label": 76, "increas": [76, 77, 177], "thick": 76, "convolv": 76, "non": [76, 151], "bw": 76, "bridg": 76, "gap": 76, "eros": 76, "spatial_axi": 76, "prob": [76, 121], "flip": [76, 177], "augment": [76, 173, 176, 177], "sign": 76, "gradient": 76, "chang": [76, 173, 176, 177], "axi": [76, 168], "NOT": 76, "probabl": 76, "constraint": 76, "annealing_step": 77, "5000": 77, "init_vari": 77, "anneal": [77, 108], "varianc": [77, 101], "nois": [77, 108], "real": [77, 107, 123], "fake": [77, 123], "harder": 77, "distribut": [77, 102], "overlap": 77, "nice": 77, "theoret": 77, "implic": 77, "qualiti": 77, "also": [77, 101, 138, 173, 176, 177], "curriculum": 77, "learn": [77, 83, 84, 177], "techniqu": 77, "iter": 77, "unblur": 77, "over": 77, "linearli": 77, "activ": [79, 107, 115, 138, 143, 145, 174], "ident": [79, 112, 115, 138, 173, 176], "_type": 79, "_dtype_lik": 79, "_supportsdtyp": 79, "supportsindex": 79, "_dtypedict": 79, "uint8": 79, "rescale_dtyp": 79, "purpos": 79, "rescal": [79, 173, 176], "dtypelik": 79, "multichannel": 81, "concaten": 81, "cuda": 81, "extract_embed": [82, 84], "get_predict_mask": [82, 84], "get_context_embed": [82, 85], "get_mask": [82, 85, 146, 148], "get_target_embed": [82, 85], "remove_first_dim": [82, 85], "update_teach": [82, 85], "supervis": [83, 84], "featur": [83, 84, 85, 110, 113, 114, 116, 123, 138, 143, 144, 145, 150, 151, 152, 153, 154, 173, 176], "access": [83, 84, 85], "momentum": [83, 84, 85], "exponenti": [83, 84, 85], "move": [83, 84, 85], "averag": [83, 84, 85, 126], "998": [83, 84, 85], "world": 84, "translat": 84, "latent": [84, 90], "space": [84, 90], "target_domain_kei": 84, "domain": [84, 153], "predict_target": 84, "devic": [84, 155], "tmp_dir": 87, "strict": 87, "ckpt": [87, 177], "lst": 88, "o2_spharm_encod": [89, 95], "abstract_prior": [89, 99], "diagonalgaussianprior": [89, 99, 101], "isotropicgaussianprior": [89, 99, 101], "compute_tc_penalti": [89, 99, 101], "identity_prior": [89, 99], "identityprior": [89, 99, 102], "joint_prior": [89, 99], "jointprior": [89, 99, 103], "calculate_elbo": [89, 90], "calculate_rcl": [89, 90], "calculate_rcl_dict": [89, 90], "decod": [89, 90, 105, 141, 143, 152, 154], "sample_z": [89, 90], "beta": [90, 130], "kld": 90, "term": [90, 138], "reconstruction_loss": [90, 111, 112], "reconstruct": [90, 111, 112, 160, 169], "_loss": [90, 119, 121, 122, 124, 130, 131], "abstractprior": 90, "decoder_latent_part": 90, "decoer": 90, "what": [90, 167], "depend": [90, 174], "xhat": 90, "input_kei": 90, "return_param": 90, "z_parts_param": 90, "kl_diverg": [99, 100, 101, 103], "param_s": [99, 100, 101, 102, 103], "kl": [100, 101, 102, 103], "properti": [100, 101, 102, 103, 147, 148, 149, 152], "logvar": 101, "learn_mean": 101, "learn_logvar": 101, "clamp_logvar": 101, "tc_penalty_weight": 101, "mu1": 101, "mu2": 101, "logvar1": 101, "logvar2": 101, "reduct": [101, 103, 119, 121, 122, 130, 131], "sum": [101, 103, 130, 144, 151, 154], "kullback": 101, "leibler": 101, "diverg": 101, "diagon": 101, "necessarili": 101, "isotrop": 101, "wise": 101, "second": [101, 151], "doesn": [102, 174], "contribut": 102, "effect": 102, "dirac": 102, "delta": 102, "z_param": 103, "multi_scale_discrimin": [105, 106], "multiscalediscrimin": [105, 106, 107], "n_layer_discrimin": [105, 106], "nlayerdiscrimin": [105, 106, 107, 108], "base_head": [105, 109], "basehead": [105, 109, 110, 111, 113, 114, 115, 116], "gan_head": [105, 109], "ganhead": [105, 109, 111, 112], "gan_head_superr": [105, 109], "ganhead_res": [105, 109, 112], "mae_head": [105, 109], "maehead": [105, 109, 113], "mask_head": [105, 109], "maskhead": [105, 109, 114], "res_blocks_head": [105, 109], "resblockshead": [105, 109, 112, 115], "vic_reg": [105, 109, 118], "vicreghead": [105, 109, 116], "adversarial_loss": [105, 118], "adversarialloss": [105, 118, 119], "chamfer_loss": [105, 118], "chamferloss": [105, 118, 120], "continuous_bernoulli": [105, 118], "cblogloss": [105, 118, 121], "cosine_loss": [105, 118], "cosineloss": [105, 118, 122], "gan_loss": [105, 111, 112, 118], "ganloss": [105, 111, 112, 118, 123], "pix2pixhd": [105, 107, 111, 112, 118, 123], "gaussian_nll_loss": [105, 118], "gaussiannllloss": [105, 118, 124], "geomloss": [105, 118], "loss_wrapp": [105, 118], "cmap_loss": [105, 118, 126], "losswrapp": [105, 118, 126], "spharm_loss": [105, 118], "threshold_loss": [105, 118], "thresholdloss": [105, 118, 128], "vicregloss": [105, 118, 129], "off_diagon": [105, 118, 129], "weibul": [105, 118], "weibulllogloss": [105, 118, 130], "weibull_log_prob": [105, 118, 130], "weighted_mse_loss": [105, 118], "weightedmseloss": [105, 118, 131], "dgcnn": [105, 133], "folding_net": [105, 133], "graph_funct": [105, 133], "vnn": [105, 133], "crossmae_decod": [105, 141, 150], "mae_decod": [105, 141, 150], "hieraencod": [105, 141, 151], "jepaencod": [105, 141, 151], "mae_encod": [105, 141, 151], "spatialmerg": [105, 141, 151], "mae": [105, 141], "hierama": [105, 141, 152], "mae_bas": [105, 141, 152], "iwmpredictor": [105, 141, 153], "jepapredictor": [105, 141, 153], "encodedskip": [105, 141, 154], "seg_vit": [105, 141, 154], "superresdecod": [105, 141, 154], "get_positional_embed": [105, 141, 155], "match_tuple_dimens": [105, 141, 155], "random_index": [105, 141, 155], "take_index": [105, 141, 155], "n_scale": 107, "input_nc": 108, "ndf": 108, "n_layer": 108, "norm_lay": [108, 143, 144, 145], "instancenorm": 108, "instancenorm3d": 108, "defin": [108, 123], "patchgan": 108, "n_channel": 108, "input_im": 108, "model_output": 108, "conv_fn": 108, "later": 108, "requires_featur": 108, "standard": 108, "generate_io_map": [109, 110], "update_param": [109, 110], "input_filenam": 110, "backbone_featur": [110, 111, 113, 114, 116], "y_hat": [110, 111, 113, 114, 116], "calcul": [110, 113, 114, 116, 126], "bcewithlogitsloss": [111, 112, 123], "feature_matching_loss": [111, 112], "l1loss": [111, 112], "mseloss": [111, 112], "reconstruction_loss_weight": [111, 112], "in_channel": [112, 115, 117, 138, 177], "out_channel": [112, 115, 117, 138, 154, 177], "final_act": [112, 115], "lr": [112, 115], "n_conv": [112, 115], "dropout": [112, 115, 117, 138, 145], "upsample_method": [112, 115], "pixelshuffl": [112, 115], "upsample_ratio": [112, 115], "first_lay": [112, 115], "dens": [112, 115], "inherit": 112, "__init__": 112, "calculate_metr": 114, "convolut": [115, 138, 148, 149, 151, 152, 154, 173, 176], "final": 115, "done": 115, "hr": 115, "length": [115, 126, 140, 155], "convprojectionlay": 115, "2048": 116, "8192": 116, "hidden_lay": [116, 132], "projectornetwork": 116, "hidden": [116, 145], "scale_factor": 117, "align_corn": 117, "antialia": 117, "batch_pairwise_dist": [118, 120], "batch_pairwise_dist2": [118, 120], "get_target_tensor": [118, 123], "get_feature_matching_loss": [118, 123], "get_gan_loss": [118, 123], "argmax": 119, "squeez": 119, "return_pr": 119, "sentinal4d": 120, "cellshap": 120, "blob": 120, "cellshape_cloud": 120, "vendor": 120, "chamfer_dist": 120, "py": [120, 177], "licens": 120, "license_antao": 120, "y": [120, 129, 143], "continu": 121, "bernoulli": 121, "propos": 121, "arxiv": 121, "ab": 121, "1907": 121, "06845": 121, "gan_mod": 123, "vanilla": 123, "target_real_label": 123, "target_fake_label": 123, "abstract": [123, 148, 152, 173, 176], "awai": 123, "need": [123, 173, 176, 177], "lsgan": 123, "wgangp": 123, "note": [123, 177], "sigmoid": [123, 173, 176], "target_is_r": 123, "fill": 123, "loss_weight": 123, "fm": 123, "10": [123, 124, 130, 138], "feature_typ": 123, "mean_dim": 124, "ep": [124, 130, 138, 168], "1e": [124, 130, 138, 168], "costmap": 126, "differenti": [126, 173, 176], "emphas": 126, "throughout": 126, "loss_fn": [126, 128], "channel_weight": 126, "_supportsarrai": 126, "_nestedsequ": 126, "complex": [126, 173, 176], "byte": 126, "loss_scal": 126, "arraylik": 126, "abov": 128, "num_featur": 129, "sim_coeff": 129, "std_coeff": 129, "cov_coeff": 129, "alpha": 130, "256": [132, 139], "scale_output": 132, "stride": [138, 145, 151, 152], "subunit": 138, "adn_ord": 138, "nda": 138, "prelu": 138, "norm": [138, 144], "dropout_dim": 138, "dilat": 138, "bia": [138, 145], "last_conv_onli": 138, "residu": 138, "conv": 138, "init": 138, "normalized_shap": 138, "print": [138, 167, 172], "sequenti": 138, "unit0": 138, "conv3d": 138, "adn": 138, "num_paramet": 138, "n": [138, 173, 174, 176], "layernorm": [138, 143, 144, 145], "05": 138, "elementwise_affin": 138, "unit1": 138, "zero": [138, 172], "entir": 138, "larger": 138, "control": [138, 172], "implicit": 138, "both": [138, 177], "side": 138, "in_filt": 139, "out_filt": 139, "n_input_ch": 139, "512": 139, "n_conv_filt": 139, "pos_embedding_length": 140, "120": 140, "emb_dim": [140, 143, 145, 147, 148, 149, 150, 151, 152, 153, 154, 155], "num_lay": [140, 144, 150, 151, 153, 154], "num_head": [140, 143, 145, 150, 151, 152, 153], "num_class": 140, "token": [140, 145, 147, 149, 150, 155], "classifi": 140, "them": [140, 173, 176], "d_model": 140, "matrix": [140, 168], "patchifi": [141, 142, 151], "cross_attent": [141, 142], "crossattent": [141, 142, 143], "crossattentionblock": [141, 142, 143], "crossselfblock": [141, 142, 143], "intermediate_weigh": [141, 142], "intermediateweigh": [141, 142, 144], "masked_unit_attent": [141, 142], "hierablock": [141, 142, 145], "maskunitattent": [141, 142, 145], "add_mask_token": [141, 150], "adjust_indices_for_cl": [141, 150], "init_weight": [141, 150, 151, 153], "init_decod": [141, 152], "init_encod": [141, 152], "predict_target_featur": [141, 153], "patchify_bas": [142, 146], "patchifybas": [142, 146, 147, 148, 149], "patchify_hiera": [142, 146], "patchifyhiera": [142, 146, 149], "take_indexes_mask": [142, 146, 149], "encoder_dim": 143, "decoder_dim": [143, 152], "qkv_bia": [143, 145], "qk_scale": 143, "attn_drop": [143, 145], "proj_drop": [143, 145], "mlp_ratio": [143, 145], "drop_path": [143, 145], "act_lay": [143, 145], "gelu": [143, 145], "in_featur": 143, "hidden_featur": 143, "out_featur": 143, "embed_dim": [144, 151], "n_output": 144, "intermedi": [144, 151, 154], "dim_out": 145, "q_stride": [145, 151, 152], "patches_per_mask_unit": 145, "attent": [145, 151, 152], "unit": [145, 149, 151, 152], "learnabl": [145, 148, 150, 151, 152, 153, 155], "create_img2token": [146, 147, 149], "extract_visible_token": [146, 147, 148, 149], "get_mask_arg": [146, 147, 148, 149], "img2token": [146, 147, 148, 149], "create_conv": [146, 148], "create_patch2img": [146, 148], "n_patch": [147, 148, 149], "context_pixel": [147, 148, 149, 151, 152], "learnable_pos_embed": [147, 148, 150, 151, 152, 153], "rearrang": 147, "forward_index": [147, 149, 150], "n_visible_patch": [147, 148, 149], "pix": [148, 149], "sin": [148, 150, 151, 152, 153, 168], "co": [148, 150, 151, 152, 153, 168], "empir": [148, 150, 151, 152, 153], "brightfield": [148, 150, 151, 152, 153], "boolean": [148, 177], "mask_units_per_dim": 149, "mask_unit": 149, "enc_dim": 150, "768": [150, 152, 154], "192": [150, 151, 152, 153, 154], "has_cls_token": 150, "inspir": 150, "crossma": [150, 152], "attend": 150, "visibl": 150, "cl": [150, 155], "backward_index": 150, "num_mask_unit": [151, 152], "architectur": [151, 152, 177], "save_lay": 151, "multihead": [151, 152], "self_attent": [151, 152], "On": 151, "doubl": 151, "pool": 151, "doe": [151, 177], "while": [151, 177], "produc": [151, 173, 176], "context": [151, 153], "n_intermediate_weight": 151, "75": [151, 152, 154], "downsample_factor": 151, "in_dim": 151, "out_dim": 151, "hiera": 151, "lowest": 151, "decoder_lay": [152, 154], "decoder_head": 152, "use_crossma": 152, "features_onli": 152, "encoder_lay": [152, 154], "encoder_head": [152, 154], "input_dim": 153, "special": 153, "condition": 153, "fluoresc": 153, "context_emb": 153, "target_mask": 153, "target_domain": 153, "n_decoder_filt": 154, "base_patch_s": 154, "upsample_factor": 154, "6134": 154, "5005": 154, "encoder_ckpt": 154, "freeze_encod": 154, "encoder_kwarg": 154, "simpl": [154, 173, 176], "aic": 154, "20x": 154, "100x": 154, "unet": 154, "like": [154, 177], "fed": [154, 173, 176], "skip": 154, "use_cls_token": 155, "ensur": [155, 177], "read_pcloud": [156, 157], "generate_img": [160, 165], "compute_rotation_matrix": [160, 168], "data_cfg": 161, "load_param": 162, "kv": 163, "item": 163, "nest": 163, "form": [163, 168], "interpol": [163, 173, 176], "limit": [164, 173, 176, 177], "fn": 164, "num_sampl": 165, "dummy_metadata": 165, "10000": 165, "mani": [165, 173, 176, 177], "k": 165, "friendli": 166, "save_to_fil": 167, "prompt": 167, "user": [167, 173, 176], "print_ord": 167, "resolv": 167, "content": [167, 172], "rich": [167, 172], "librari": 167, "tree": 167, "hydra": [167, 173, 176, 177], "refer": [167, 173, 176], "export": 167, "pad_valu": 168, "06": 168, "pose": 168, "about": 168, "equivari": [168, 174], "theta": 168, "indic": 168, "pair": 168, "precis": 168, "sure": 172, "close": 172, "properli": 172, "prevent": 172, "failur": 172, "multirun": 172, "warn": 172, "metric_dict": 172, "metric_nam": 172, "safe": 172, "callbacks_cfg": 172, "logger_cfg": 172, "object_dict": 172, "rank": 172, "task_func": 172, "decor": 172, "resist": 172, "finish": 172, "except": 172, "occur": 172, "total": 172, "execut": 172, "dir": 172, "aim": [173, 176], "common": [173, 176], "manner": [173, 176], "easi": [173, 176], "beginn": [173, 176], "flexibl": [173, 176, 177], "custom": [173, 176], "templat": [173, 176, 177], "conveni": [173, 176], "syntax": [173, 176], "repo": [173, 176], "best": [173, 176], "familiar": [173, 176, 177], "mechan": [173, 176], "topic": [173, 176], "cover": [173, 176], "becaus": [173, 176], "help": [173, 176], "These": [173, 176], "reli": [173, 176], "heavili": [173, 176], "omegaconf": [173, 176], "expos": [173, 176], "_aux": [173, 176], "section": [173, 176], "persist": [173, 176], "speed": [173, 176], "up": [173, 176, 177], "divid": [173, 176], "manifest": [173, 176], "preprocess": [173, 176], "offer": [173, 176], "wide": [173, 176], "often": [173, 176, 177], "primari": [173, 176], "due": [173, 176], "score": [173, 176], "v": [173, 176], "binar": [173, 176, 177], "own": [173, 176], "thing": [173, 176], "pin_memori": [173, 176], "workhors": [173, 176], "source_col": [173, 176, 177], "target_col": [173, 176], "organ": [173, 176], "arrang": [173, 176], "through": [173, 176], "further": [173, 176], "lr_schedul": [173, 176], "alwai": [173, 176], "well": [173, 176], "_target_": [173, 176], "_partial_": [173, 176], "sigmoid_rescal": [173, 176], "tell": [173, 176], "bit": [173, 176], "encourag": [173, 176], "few": [173, 176, 177], "git": 174, "clone": 174, "allencellmodel": 174, "cd": 174, "conda": 174, "environ": 174, "myenv": 174, "pip": 174, "txt": 174, "relat": 174, "equiv": 174, "upgrad": 174, "subpackag": 175, "outlin": 177, "modif": 177, "suit": 177, "aicsimageio": 177, "recommend": 177, "suitabl": 177, "robust": 177, "again": 177, "excel": 177, "naiv": 177, "incorrect": 177, "neural": 177, "signal": 177, "cut": 177, "crucial": 177, "bottleneck": 177, "suggest": 177, "look": 177, "consider": 177, "primarili": 177, "As": 177, "thumb": 177, "enough": 177, "entiti": 177, "establish": 177, "too": 177, "downsampl": 177, "decreas": 177, "usag": 177, "ckpt_path": 177, "dictat": 177, "omit": 177}, "objects": {"": [[1, 0, 0, "-", "cyto_dl"]], "cyto_dl": [[2, 0, 0, "-", "api"], [8, 0, 0, "-", "callbacks"], [17, 0, 0, "-", "dataframe"], [24, 0, 0, "-", "datamodules"], [35, 0, 0, "-", "eval"], [36, 0, 0, "-", "image"], [58, 0, 0, "-", "loggers"], [60, 0, 0, "-", "models"], [105, 0, 0, "-", "nn"], [156, 0, 0, "-", "point_cloud"], [159, 0, 0, "-", "train"], [160, 0, 0, "-", "utils"]], "cyto_dl.api": [[3, 0, 0, "-", "cyto_dl_model"], [6, 0, 0, "-", "data"], [7, 0, 0, "-", "model"]], "cyto_dl.api.cyto_dl_model": [[4, 0, 0, "-", "cyto_dl_base_model"], [5, 0, 0, "-", "segmentation_plugin_model"]], "cyto_dl.api.cyto_dl_model.cyto_dl_base_model": [[4, 1, 1, "", "CytoDLBaseModel"]], "cyto_dl.api.cyto_dl_model.cyto_dl_base_model.CytoDLBaseModel": [[4, 2, 1, "", "from_default_config"], [4, 2, 1, "", "from_existing_config"], [4, 2, 1, "", "get_config"], [4, 2, 1, "", "get_experiment_name"], [4, 2, 1, "", "predict"], [4, 2, 1, "", "save_config"], [4, 2, 1, "", "set_experiment_name"], [4, 2, 1, "", "train"]], "cyto_dl.api.cyto_dl_model.segmentation_plugin_model": [[5, 1, 1, "", "SegmentationPluginModel"]], "cyto_dl.api.cyto_dl_model.segmentation_plugin_model.SegmentationPluginModel": [[5, 2, 1, "", "get_hardware_type"], [5, 2, 1, "", "get_input_channel"], [5, 2, 1, "", "get_manifest_column_names"], [5, 2, 1, "", "get_patch_size"], [5, 2, 1, "", "get_raw_image_channels"], [5, 2, 1, "", "get_split_column"], [5, 2, 1, "", "remove_split_column"], [5, 2, 1, "", "set_hardware_type"], [5, 2, 1, "", "set_input_channel"], [5, 2, 1, "", "set_manifest_column_names"], [5, 2, 1, "", "set_patch_size"], [5, 2, 1, "", "set_raw_image_channels"], [5, 2, 1, "", "set_split_column"]], "cyto_dl.api.data": [[6, 1, 1, "", "AutoThresholdMethod"], [6, 1, 1, "", "ExperimentType"], [6, 1, 1, "", "HardwareType"], [6, 1, 1, "", "PatchSize"]], "cyto_dl.api.data.AutoThresholdMethod": [[6, 3, 1, "", "ISODATA"], [6, 3, 1, "", "LI"], [6, 3, 1, "", "LOCAL"], [6, 3, 1, "", "MEAN"], [6, 3, 1, "", "MINIMUM"], [6, 3, 1, "", "MULTIOTSU"], [6, 3, 1, "", "NIBLACK"], [6, 3, 1, "", "OTSU"], [6, 3, 1, "", "SAUVOLA"], [6, 3, 1, "", "TRIANGLE"], [6, 3, 1, "", "YEN"]], "cyto_dl.api.data.ExperimentType": [[6, 3, 1, "", "GAN"], [6, 3, 1, "", "INSTANCE_SEG"], [6, 3, 1, "", "LABEL_FREE"], [6, 3, 1, "", "SEGMENTATION"], [6, 3, 1, "", "SEGMENTATION_ARRAY"], [6, 3, 1, "", "SEGMENTATION_PLUGIN"]], "cyto_dl.api.data.HardwareType": [[6, 3, 1, "", "CPU"], [6, 3, 1, "", "GPU"]], "cyto_dl.api.data.PatchSize": [[6, 3, 1, "", "LARGE"], [6, 3, 1, "", "MEDIUM"], [6, 3, 1, "", "SMALL"]], "cyto_dl.api.model": [[7, 1, 1, "", "CytoDLModel"]], "cyto_dl.api.model.CytoDLModel": [[7, 2, 1, "", "download_example_data"], [7, 2, 1, "", "load_config_from_dict"], [7, 2, 1, "", "load_config_from_file"], [7, 2, 1, "", "load_default_experiment"], [7, 2, 1, "", "override_config"], [7, 2, 1, "", "predict"], [7, 2, 1, "", "print_config"], [7, 2, 1, "", "save_config"], [7, 2, 1, "", "train"]], "cyto_dl.callbacks": [[10, 0, 0, "-", "csv_saver"], [11, 0, 0, "-", "image_saver"], [13, 0, 0, "-", "layer_freeze"], [14, 0, 0, "-", "model_utils"], [15, 0, 0, "-", "outlier_detection"]], "cyto_dl.callbacks.csv_saver": [[10, 1, 1, "", "CSVSaver"], [10, 1, 1, "", "JEPASaver"]], "cyto_dl.callbacks.csv_saver.CSVSaver": [[10, 2, 1, "", "on_predict_epoch_end"], [10, 2, 1, "", "pred_to_csv"], [10, 2, 1, "", "save_feats"]], "cyto_dl.callbacks.csv_saver.JEPASaver": [[10, 2, 1, "", "pred_to_csv"]], "cyto_dl.callbacks.image_saver": [[11, 1, 1, "", "ImageSaver"]], "cyto_dl.callbacks.image_saver.ImageSaver": [[11, 2, 1, "", "on_predict_batch_end"], [11, 2, 1, "", "on_test_batch_end"], [11, 2, 1, "", "on_train_batch_end"], [11, 2, 1, "", "on_validation_batch_end"], [11, 2, 1, "", "save"]], "cyto_dl.callbacks.layer_freeze": [[13, 1, 1, "", "LayerFreeze"]], "cyto_dl.callbacks.layer_freeze.LayerFreeze": [[13, 2, 1, "", "setup"]], "cyto_dl.callbacks.model_utils": [[14, 1, 1, "", "GetEmbeddings"], [14, 4, 1, "", "get_all_embeddings"], [14, 4, 1, "", "save_predictions_classifier"]], "cyto_dl.callbacks.model_utils.GetEmbeddings": [[14, 2, 1, "", "on_test_epoch_end"]], "cyto_dl.callbacks.outlier_detection": [[15, 1, 1, "", "OutlierDetection"]], "cyto_dl.callbacks.outlier_detection.OutlierDetection": [[15, 2, 1, "", "calculate_mahalanobis_hook"], [15, 2, 1, "", "flatten_activations"], [15, 2, 1, "", "on_load_checkpoint"], [15, 2, 1, "", "on_predict_batch_end"], [15, 2, 1, "", "on_predict_epoch_start"], [15, 2, 1, "", "on_save_checkpoint"], [15, 2, 1, "", "on_test_batch_end"], [15, 2, 1, "", "on_test_epoch_start"], [15, 2, 1, "", "on_train_epoch_start"], [15, 2, 1, "", "update_covariance_hook"]], "cyto_dl.dataframe": [[18, 0, 0, "-", "readers"], [19, 0, 0, "-", "transforms"]], "cyto_dl.dataframe.readers": [[18, 4, 1, "", "filter_columns"], [18, 4, 1, "", "read_csv"], [18, 4, 1, "", "read_dataframe"], [18, 4, 1, "", "read_h5ad"], [18, 4, 1, "", "read_parquet"]], "cyto_dl.dataframe.transforms": [[20, 0, 0, "-", "filter"], [21, 0, 0, "-", "group_cols"], [22, 0, 0, "-", "misc"], [23, 0, 0, "-", "split"]], "cyto_dl.dataframe.transforms.filter": [[20, 4, 1, "", "filter_columns"], [20, 4, 1, "", "filter_rows"]], "cyto_dl.dataframe.transforms.group_cols": [[21, 1, 1, "", "GroupCols"]], "cyto_dl.dataframe.transforms.misc": [[22, 4, 1, "", "append_class_weights"], [22, 4, 1, "", "append_labels_to_integers"], [22, 4, 1, "", "append_one_hot"], [22, 4, 1, "", "make_random_df"]], "cyto_dl.dataframe.transforms.split": [[23, 4, 1, "", "sample_n_each"], [23, 4, 1, "", "split_dataframe"]], "cyto_dl.datamodules": [[25, 0, 0, "-", "array"], [26, 0, 0, "-", "data_dict"], [27, 0, 0, "-", "dataframe"], [31, 0, 0, "-", "folder"], [32, 0, 0, "-", "multidim_image"], [33, 0, 0, "-", "smartcache"], [34, 0, 0, "-", "torchvision"]], "cyto_dl.datamodules.array": [[25, 4, 1, "", "make_array_dataloader"]], "cyto_dl.datamodules.data_dict": [[26, 4, 1, "", "make_data_dict_dataloader"]], "cyto_dl.datamodules.dataframe": [[28, 0, 0, "-", "dataframe_datamodule"], [29, 0, 0, "-", "grouped_dataframe_datamodule"], [30, 0, 0, "-", "utils"]], "cyto_dl.datamodules.dataframe.dataframe_datamodule": [[28, 1, 1, "", "DataframeDatamodule"]], "cyto_dl.datamodules.dataframe.dataframe_datamodule.DataframeDatamodule": [[28, 2, 1, "", "get_dataloader"], [28, 2, 1, "", "get_dataset"], [28, 2, 1, "", "make_dataloader"], [28, 2, 1, "", "predict_dataloader"], [28, 2, 1, "", "test_dataloader"], [28, 2, 1, "", "train_dataloader"], [28, 2, 1, "", "val_dataloader"]], "cyto_dl.datamodules.dataframe.grouped_dataframe_datamodule": [[29, 1, 1, "", "GroupedDataframeDatamodule"]], "cyto_dl.datamodules.dataframe.grouped_dataframe_datamodule.GroupedDataframeDatamodule": [[29, 2, 1, "", "make_dataloader"]], "cyto_dl.datamodules.dataframe.utils": [[30, 1, 1, "", "AlternatingBatchSampler"], [30, 1, 1, "", "RemoveNaNKeysd"], [30, 4, 1, "", "get_canonical_split_name"], [30, 4, 1, "", "get_dataset"], [30, 4, 1, "", "make_multiple_dataframe_splits"], [30, 4, 1, "", "make_single_dataframe_splits"], [30, 4, 1, "", "parse_transforms"]], "cyto_dl.datamodules.folder": [[31, 4, 1, "", "make_folder_dataloader"]], "cyto_dl.datamodules.multidim_image": [[32, 1, 1, "", "MultiDimImageDataset"], [32, 4, 1, "", "make_multidim_image_dataloader"]], "cyto_dl.datamodules.multidim_image.MultiDimImageDataset": [[32, 2, 1, "", "create_metatensor"], [32, 2, 1, "", "get_per_file_args"], [32, 2, 1, "", "is_batch"]], "cyto_dl.datamodules.smartcache": [[33, 1, 1, "", "SmartcacheDatamodule"]], "cyto_dl.datamodules.smartcache.SmartcacheDatamodule": [[33, 2, 1, "", "get_per_file_args"], [33, 2, 1, "", "make_dataloader"], [33, 2, 1, "", "predict_dataloader"], [33, 2, 1, "", "prepare_data"], [33, 2, 1, "", "setup"], [33, 2, 1, "", "test_dataloader"], [33, 2, 1, "", "train_dataloader"], [33, 2, 1, "", "val_dataloader"]], "cyto_dl.datamodules.torchvision": [[34, 1, 1, "", "CIFAR10"], [34, 1, 1, "", "CIFAR100"], [34, 1, 1, "", "CLEVRClassification"], [34, 1, 1, "", "Caltech101"], [34, 1, 1, "", "Caltech256"], [34, 1, 1, "", "CelebA"], [34, 1, 1, "", "Cityscapes"], [34, 1, 1, "", "CocoCaptions"], [34, 1, 1, "", "CocoDetection"], [34, 1, 1, "", "Country211"], [34, 1, 1, "", "DTD"], [34, 1, 1, "", "DatasetFolder"], [34, 1, 1, "", "EMNIST"], [34, 1, 1, "", "EuroSAT"], [34, 1, 1, "", "FER2013"], [34, 1, 1, "", "FGVCAircraft"], [34, 1, 1, "", "FakeData"], [34, 1, 1, "", "FashionMNIST"], [34, 1, 1, "", "Flickr30k"], [34, 1, 1, "", "Flickr8k"], [34, 1, 1, "", "Flowers102"], [34, 1, 1, "", "Food101"], [34, 1, 1, "", "GTSRB"], [34, 1, 1, "", "HMDB51"], [34, 1, 1, "", "INaturalist"], [34, 1, 1, "", "ImageFolder"], [34, 1, 1, "", "ImageNet"], [34, 1, 1, "", "KMNIST"], [34, 1, 1, "", "Kinetics"], [34, 1, 1, "", "Kitti"], [34, 1, 1, "", "LFWPairs"], [34, 1, 1, "", "LFWPeople"], [34, 1, 1, "", "LSUN"], [34, 1, 1, "", "LSUNClass"], [34, 1, 1, "", "MNIST"], [34, 1, 1, "", "MovingMNIST"], [34, 1, 1, "", "Omniglot"], [34, 1, 1, "", "OxfordIIITPet"], [34, 1, 1, "", "PCAM"], [34, 1, 1, "", "PhotoTour"], [34, 1, 1, "", "Places365"], [34, 1, 1, "", "QMNIST"], [34, 1, 1, "", "RenderedSST2"], [34, 1, 1, "", "SBDataset"], [34, 1, 1, "", "SBU"], [34, 1, 1, "", "SEMEION"], [34, 1, 1, "", "STL10"], [34, 1, 1, "", "SUN397"], [34, 1, 1, "", "SVHN"], [34, 1, 1, "", "StanfordCars"], [34, 1, 1, "", "UCF101"], [34, 1, 1, "", "USPS"], [34, 1, 1, "", "VOCDetection"], [34, 1, 1, "", "VOCSegmentation"], [34, 1, 1, "", "VisionDataset"], [34, 1, 1, "", "WIDERFace"]], "cyto_dl.eval": [[35, 4, 1, "", "main"]], "cyto_dl.image": [[37, 0, 0, "-", "io"], [44, 0, 0, "-", "transforms"]], "cyto_dl.image.io": [[38, 0, 0, "-", "aicsimage_loader"], [39, 0, 0, "-", "monai_bio_reader"], [40, 0, 0, "-", "numpy_reader"], [41, 0, 0, "-", "ome_zarr_reader"], [42, 0, 0, "-", "polygon_loader"], [43, 0, 0, "-", "skimage_reader"]], "cyto_dl.image.io.aicsimage_loader": [[38, 1, 1, "", "AICSImageLoaderd"]], "cyto_dl.image.io.aicsimage_loader.AICSImageLoaderd": [[38, 2, 1, "", "split_args"]], "cyto_dl.image.io.monai_bio_reader": [[39, 1, 1, "", "MonaiBioReader"]], "cyto_dl.image.io.monai_bio_reader.MonaiBioReader": [[39, 2, 1, "", "get_data"], [39, 2, 1, "", "read"], [39, 2, 1, "", "verify_suffix"]], "cyto_dl.image.io.numpy_reader": [[40, 1, 1, "", "ReadNumpyFile"]], "cyto_dl.image.io.ome_zarr_reader": [[41, 1, 1, "", "OmeZarrReader"]], "cyto_dl.image.io.ome_zarr_reader.OmeZarrReader": [[41, 2, 1, "", "get_data"], [41, 2, 1, "", "read"], [41, 2, 1, "", "verify_suffix"]], "cyto_dl.image.io.polygon_loader": [[42, 1, 1, "", "PolygonLoaderd"]], "cyto_dl.image.io.skimage_reader": [[43, 1, 1, "", "SkimageReader"]], "cyto_dl.image.io.skimage_reader.SkimageReader": [[43, 2, 1, "", "get_data"], [43, 2, 1, "", "read"], [43, 2, 1, "", "verify_suffix"]], "cyto_dl.image.transforms": [[45, 0, 0, "-", "add_meta"], [46, 0, 0, "-", "bright_sampler"], [47, 0, 0, "-", "clip"], [48, 0, 0, "-", "contrastadjust"], [49, 0, 0, "-", "generate_jepa_masks"], [50, 0, 0, "-", "merge"], [51, 0, 0, "-", "multiscale_cropper"], [52, 0, 0, "-", "pad"], [53, 0, 0, "-", "project"], [55, 0, 0, "-", "save"], [57, 0, 0, "-", "track_transforms"]], "cyto_dl.image.transforms.add_meta": [[45, 1, 1, "", "AddMeta"]], "cyto_dl.image.transforms.bright_sampler": [[46, 1, 1, "", "BrightSampler"]], "cyto_dl.image.transforms.clip": [[47, 1, 1, "", "Clip"], [47, 1, 1, "", "Clipd"]], "cyto_dl.image.transforms.contrastadjust": [[48, 1, 1, "", "ContrastAdjust"], [48, 1, 1, "", "ContrastAdjustd"]], "cyto_dl.image.transforms.generate_jepa_masks": [[49, 1, 1, "", "JEPAMaskGenerator"]], "cyto_dl.image.transforms.generate_jepa_masks.JEPAMaskGenerator": [[49, 2, 1, "", "remove_excess_pixels"]], "cyto_dl.image.transforms.merge": [[50, 1, 1, "", "Merged"]], "cyto_dl.image.transforms.multiscale_cropper": [[51, 1, 1, "", "RandomMultiScaleCropd"]], "cyto_dl.image.transforms.multiscale_cropper.RandomMultiScaleCropd": [[51, 2, 1, "", "generate_slices"]], "cyto_dl.image.transforms.pad": [[52, 1, 1, "", "PadZd"]], "cyto_dl.image.transforms.project": [[53, 1, 1, "", "MaxProjectd"]], "cyto_dl.image.transforms.save": [[55, 1, 1, "", "Save"], [55, 1, 1, "", "Saved"]], "cyto_dl.image.transforms.track_transforms": [[57, 1, 1, "", "CropResize"], [57, 1, 1, "", "GenerateTrackLabels"], [57, 1, 1, "", "PerChannel"], [57, 1, 1, "", "SplitTrackd"]], "cyto_dl.loggers": [[59, 0, 0, "-", "mlflow"]], "cyto_dl.loggers.mlflow": [[59, 1, 1, "", "MLFlowLogger"]], "cyto_dl.loggers.mlflow.MLFlowLogger": [[59, 2, 1, "", "after_save_checkpoint"], [59, 2, 1, "", "log_hyperparams"], [59, 2, 1, "", "log_metrics"]], "cyto_dl.models": [[61, 0, 0, "-", "base_model"], [62, 0, 0, "-", "basic_model"], [63, 0, 0, "-", "classification"], [66, 0, 0, "-", "contrastive"], [72, 0, 0, "-", "im2im"], [82, 0, 0, "-", "jepa"], [86, 0, 0, "-", "utils"]], "cyto_dl.models.base_model": [[61, 1, 1, "", "BaseModel"], [61, 1, 1, "", "BaseModelMeta"]], "cyto_dl.models.base_model.BaseModel": [[61, 2, 1, "", "compute_metrics"], [61, 2, 1, "", "configure_optimizers"], [61, 2, 1, "", "forward"], [61, 2, 1, "", "model_step"], [61, 2, 1, "", "on_train_start"], [61, 2, 1, "", "parse_batch"], [61, 2, 1, "", "predict_step"], [61, 2, 1, "", "test_step"], [61, 2, 1, "", "training_step"], [61, 2, 1, "", "validation_step"]], "cyto_dl.models.basic_model": [[62, 1, 1, "", "BasicModel"]], "cyto_dl.models.basic_model.BasicModel": [[62, 2, 1, "", "forward"], [62, 2, 1, "", "model_step"]], "cyto_dl.models.classification": [[64, 0, 0, "-", "classification"], [65, 0, 0, "-", "timepoint_classification"]], "cyto_dl.models.classification.classification": [[64, 1, 1, "", "Classification"]], "cyto_dl.models.classification.classification.Classification": [[64, 2, 1, "", "forward"], [64, 2, 1, "", "model_step"], [64, 2, 1, "", "predict_step"], [64, 2, 1, "", "save_images"], [64, 2, 1, "", "should_save_image"]], "cyto_dl.models.classification.timepoint_classification": [[65, 1, 1, "", "TimepointClassification"]], "cyto_dl.models.classification.timepoint_classification.TimepointClassification": [[65, 2, 1, "", "predict_step"]], "cyto_dl.models.contrastive": [[67, 0, 0, "-", "contrastive"]], "cyto_dl.models.contrastive.contrastive": [[67, 1, 1, "", "Contrastive"]], "cyto_dl.models.contrastive.contrastive.Contrastive": [[67, 2, 1, "", "forward"], [67, 2, 1, "", "model_step"], [67, 2, 1, "", "plot_classes"], [67, 2, 1, "", "plot_neighbors"], [67, 2, 1, "", "predict_step"]], "cyto_dl.models.im2im": [[73, 0, 0, "-", "gan"], [74, 0, 0, "-", "multi_task"], [75, 0, 0, "-", "utils"]], "cyto_dl.models.im2im.gan": [[73, 1, 1, "", "GAN"]], "cyto_dl.models.im2im.gan.GAN": [[73, 2, 1, "", "configure_optimizers"], [73, 2, 1, "", "model_step"], [73, 2, 1, "", "predict_step"]], "cyto_dl.models.im2im.multi_task": [[74, 1, 1, "", "MultiTaskIm2Im"]], "cyto_dl.models.im2im.multi_task.MultiTaskIm2Im": [[74, 2, 1, "", "configure_optimizers"], [74, 2, 1, "", "forward"], [74, 2, 1, "", "get_n_postprocess_image"], [74, 2, 1, "", "get_per_head"], [74, 2, 1, "", "model_step"], [74, 2, 1, "", "predict_step"], [74, 2, 1, "", "run_forward"], [74, 2, 1, "", "test_step"], [74, 2, 1, "", "training_step"], [74, 2, 1, "", "validation_step"]], "cyto_dl.models.im2im.utils": [[76, 0, 0, "-", "instance_seg"], [77, 0, 0, "-", "noise_annealer"], [78, 0, 0, "-", "postprocessing"]], "cyto_dl.models.im2im.utils.instance_seg": [[76, 1, 1, "", "InstanceSegCluster"], [76, 1, 1, "", "InstanceSegLoss"], [76, 1, 1, "", "InstanceSegPreprocessd"], [76, 1, 1, "", "InstanceSegRandFlipd"], [76, 4, 1, "", "pad_slice"]], "cyto_dl.models.im2im.utils.instance_seg.InstanceSegCluster": [[76, 2, 1, "", "cluster_object"], [76, 2, 1, "", "kd_clustering"], [76, 2, 1, "", "remove_small_skeletons"]], "cyto_dl.models.im2im.utils.instance_seg.InstanceSegPreprocessd": [[76, 2, 1, "", "embed_from_skel"], [76, 2, 1, "", "keep_largest_cc"], [76, 2, 1, "", "label_2d"], [76, 2, 1, "", "shrink"], [76, 2, 1, "", "skeleton_tall"], [76, 2, 1, "", "smooth_embedding"], [76, 2, 1, "", "topology_preserving_thinning"]], "cyto_dl.models.im2im.utils.noise_annealer": [[77, 1, 1, "", "NoiseAnnealer"]], "cyto_dl.models.im2im.utils.noise_annealer.NoiseAnnealer": [[77, 2, 1, "", "update_noise"]], "cyto_dl.models.im2im.utils.postprocessing": [[79, 0, 0, "-", "act_thresh_label"], [80, 0, 0, "-", "auto_thresh"], [81, 0, 0, "-", "dict_to_im"]], "cyto_dl.models.im2im.utils.postprocessing.act_thresh_label": [[79, 1, 1, "", "ActThreshLabel"]], "cyto_dl.models.im2im.utils.postprocessing.auto_thresh": [[80, 1, 1, "", "AutoThreshold"]], "cyto_dl.models.im2im.utils.postprocessing.dict_to_im": [[81, 1, 1, "", "DictToIm"], [81, 4, 1, "", "detach"]], "cyto_dl.models.jepa": [[83, 0, 0, "-", "ijepa"], [84, 0, 0, "-", "iwm"], [85, 0, 0, "-", "jepa_base"]], "cyto_dl.models.jepa.ijepa": [[83, 1, 1, "", "IJEPA"]], "cyto_dl.models.jepa.ijepa.IJEPA": [[83, 2, 1, "", "model_step"]], "cyto_dl.models.jepa.iwm": [[84, 1, 1, "", "IWM"]], "cyto_dl.models.jepa.iwm.IWM": [[84, 2, 1, "", "extract_embeddings"], [84, 2, 1, "", "get_predict_masks"], [84, 2, 1, "", "model_step"], [84, 2, 1, "", "predict_step"], [84, 2, 1, "", "test_step"]], "cyto_dl.models.jepa.jepa_base": [[85, 1, 1, "", "JEPABase"]], "cyto_dl.models.jepa.jepa_base.JEPABase": [[85, 2, 1, "", "configure_optimizers"], [85, 2, 1, "", "forward"], [85, 2, 1, "", "get_context_embeddings"], [85, 2, 1, "", "get_mask"], [85, 2, 1, "", "get_target_embeddings"], [85, 2, 1, "", "model_step"], [85, 2, 1, "", "predict_step"], [85, 2, 1, "", "remove_first_dim"], [85, 2, 1, "", "update_teacher"]], "cyto_dl.models.utils": [[87, 0, 0, "-", "mlflow"], [88, 0, 0, "-", "utils"]], "cyto_dl.models.utils.mlflow": [[87, 4, 1, "", "get_config"], [87, 4, 1, "", "load_model_from_checkpoint"]], "cyto_dl.models.utils.utils": [[88, 4, 1, "", "find_indices"]], "cyto_dl.models.vae": [[90, 0, 0, "-", "base_vae"], [99, 0, 0, "-", "priors"]], "cyto_dl.models.vae.base_vae": [[90, 1, 1, "", "BaseVAE"]], "cyto_dl.models.vae.base_vae.BaseVAE": [[90, 2, 1, "", "calculate_elbo"], [90, 2, 1, "", "calculate_rcl"], [90, 2, 1, "", "calculate_rcl_dict"], [90, 2, 1, "", "decode"], [90, 2, 1, "", "encode"], [90, 2, 1, "", "forward"], [90, 2, 1, "", "model_step"], [90, 2, 1, "", "sample_z"]], "cyto_dl.models.vae.priors": [[100, 0, 0, "-", "abstract_prior"], [101, 0, 0, "-", "gaussian"], [102, 0, 0, "-", "identity_prior"], [103, 0, 0, "-", "joint_prior"]], "cyto_dl.models.vae.priors.abstract_prior": [[100, 1, 1, "", "Prior"]], "cyto_dl.models.vae.priors.abstract_prior.Prior": [[100, 2, 1, "", "forward"], [100, 2, 1, "", "kl_divergence"], [100, 5, 1, "", "param_size"], [100, 2, 1, "", "sample"]], "cyto_dl.models.vae.priors.gaussian": [[101, 1, 1, "", "DiagonalGaussianPrior"], [101, 1, 1, "", "IsotropicGaussianPrior"], [101, 4, 1, "", "compute_tc_penalty"]], "cyto_dl.models.vae.priors.gaussian.DiagonalGaussianPrior": [[101, 2, 1, "", "forward"], [101, 2, 1, "", "kl_divergence"], [101, 5, 1, "", "param_size"]], "cyto_dl.models.vae.priors.gaussian.IsotropicGaussianPrior": [[101, 2, 1, "", "forward"], [101, 2, 1, "", "kl_divergence"], [101, 5, 1, "", "param_size"], [101, 2, 1, "", "sample"]], "cyto_dl.models.vae.priors.identity_prior": [[102, 1, 1, "", "IdentityPrior"]], "cyto_dl.models.vae.priors.identity_prior.IdentityPrior": [[102, 2, 1, "", "forward"], [102, 5, 1, "", "param_size"]], "cyto_dl.models.vae.priors.joint_prior": [[103, 1, 1, "", "JointPrior"]], "cyto_dl.models.vae.priors.joint_prior.JointPrior": [[103, 2, 1, "", "forward"], [103, 2, 1, "", "kl_divergence"], [103, 5, 1, "", "param_size"], [103, 2, 1, "", "sample"]], "cyto_dl.nn": [[106, 0, 0, "-", "discriminators"], [109, 0, 0, "-", "head"], [117, 0, 0, "-", "hr_skip"], [118, 0, 0, "-", "losses"], [132, 0, 0, "-", "mlp"], [138, 0, 0, "-", "res_unit"], [139, 0, 0, "-", "spatial_transformer"], [140, 0, 0, "-", "track_sequence_predictor"], [141, 0, 0, "-", "vits"]], "cyto_dl.nn.discriminators": [[107, 0, 0, "-", "multi_scale_discriminator"], [108, 0, 0, "-", "n_layer_discriminator"]], "cyto_dl.nn.discriminators.multi_scale_discriminator": [[107, 1, 1, "", "MultiScaleDiscriminator"]], "cyto_dl.nn.discriminators.multi_scale_discriminator.MultiScaleDiscriminator": [[107, 2, 1, "", "forward"]], "cyto_dl.nn.discriminators.n_layer_discriminator": [[108, 1, 1, "", "NLayerDiscriminator"]], "cyto_dl.nn.discriminators.n_layer_discriminator.NLayerDiscriminator": [[108, 2, 1, "", "forward"]], "cyto_dl.nn.head": [[110, 0, 0, "-", "base_head"], [111, 0, 0, "-", "gan_head"], [112, 0, 0, "-", "gan_head_superres"], [113, 0, 0, "-", "mae_head"], [114, 0, 0, "-", "mask_head"], [115, 0, 0, "-", "res_blocks_head"], [116, 0, 0, "-", "vic_reg"]], "cyto_dl.nn.head.base_head": [[110, 1, 1, "", "BaseHead"]], "cyto_dl.nn.head.base_head.BaseHead": [[110, 2, 1, "", "forward"], [110, 2, 1, "", "generate_io_map"], [110, 2, 1, "", "run_head"], [110, 2, 1, "", "update_params"]], "cyto_dl.nn.head.gan_head": [[111, 1, 1, "", "GANHead"]], "cyto_dl.nn.head.gan_head.GANHead": [[111, 2, 1, "", "forward"], [111, 2, 1, "", "run_head"]], "cyto_dl.nn.head.gan_head_superres": [[112, 1, 1, "", "GANHead_resize"]], "cyto_dl.nn.head.gan_head_superres.GANHead_resize": [[112, 2, 1, "", "forward"]], "cyto_dl.nn.head.mae_head": [[113, 1, 1, "", "MAEHead"]], "cyto_dl.nn.head.mae_head.MAEHead": [[113, 2, 1, "", "run_head"]], "cyto_dl.nn.head.mask_head": [[114, 1, 1, "", "MaskHead"]], "cyto_dl.nn.head.mask_head.MaskHead": [[114, 2, 1, "", "run_head"]], "cyto_dl.nn.head.res_blocks_head": [[115, 1, 1, "", "ResBlocksHead"]], "cyto_dl.nn.head.res_blocks_head.ResBlocksHead": [[115, 2, 1, "", "forward"]], "cyto_dl.nn.head.vic_reg": [[116, 1, 1, "", "VICRegHead"]], "cyto_dl.nn.head.vic_reg.VICRegHead": [[116, 2, 1, "", "run_head"]], "cyto_dl.nn.hr_skip": [[117, 1, 1, "", "HRSkip"]], "cyto_dl.nn.hr_skip.HRSkip": [[117, 2, 1, "", "forward"]], "cyto_dl.nn.losses": [[119, 0, 0, "-", "adversarial_loss"], [120, 0, 0, "-", "chamfer_loss"], [121, 0, 0, "-", "continuous_bernoulli"], [122, 0, 0, "-", "cosine_loss"], [123, 0, 0, "-", "gan_loss"], [124, 0, 0, "-", "gaussian_nll_loss"], [126, 0, 0, "-", "loss_wrapper"], [128, 0, 0, "-", "threshold_loss"], [129, 0, 0, "-", "vic_reg"], [130, 0, 0, "-", "weibull"], [131, 0, 0, "-", "weighted_mse_loss"]], "cyto_dl.nn.losses.adversarial_loss": [[119, 1, 1, "", "AdversarialLoss"]], "cyto_dl.nn.losses.adversarial_loss.AdversarialLoss": [[119, 2, 1, "", "forward"]], "cyto_dl.nn.losses.chamfer_loss": [[120, 1, 1, "", "ChamferLoss"]], "cyto_dl.nn.losses.chamfer_loss.ChamferLoss": [[120, 2, 1, "", "batch_pairwise_dist"], [120, 2, 1, "", "batch_pairwise_dist2"], [120, 2, 1, "", "forward"]], "cyto_dl.nn.losses.continuous_bernoulli": [[121, 1, 1, "", "CBLogLoss"]], "cyto_dl.nn.losses.continuous_bernoulli.CBLogLoss": [[121, 2, 1, "", "forward"]], "cyto_dl.nn.losses.cosine_loss": [[122, 1, 1, "", "CosineLoss"]], "cyto_dl.nn.losses.cosine_loss.CosineLoss": [[122, 2, 1, "", "forward"]], "cyto_dl.nn.losses.gan_loss": [[123, 1, 1, "", "GANLoss"], [123, 1, 1, "", "Pix2PixHD"]], "cyto_dl.nn.losses.gan_loss.GANLoss": [[123, 2, 1, "", "get_target_tensor"]], "cyto_dl.nn.losses.gan_loss.Pix2PixHD": [[123, 2, 1, "", "get_feature_matching_loss"], [123, 2, 1, "", "get_gan_loss"]], "cyto_dl.nn.losses.gaussian_nll_loss": [[124, 1, 1, "", "GaussianNLLLoss"]], "cyto_dl.nn.losses.gaussian_nll_loss.GaussianNLLLoss": [[124, 2, 1, "", "forward"]], "cyto_dl.nn.losses.loss_wrapper": [[126, 1, 1, "", "CMAP_loss"], [126, 1, 1, "", "LossWrapper"]], "cyto_dl.nn.losses.threshold_loss": [[128, 1, 1, "", "ThresholdLoss"]], "cyto_dl.nn.losses.vic_reg": [[129, 1, 1, "", "VICRegLoss"], [129, 4, 1, "", "off_diagonal"]], "cyto_dl.nn.losses.vic_reg.VICRegLoss": [[129, 2, 1, "", "forward"]], "cyto_dl.nn.losses.weibull": [[130, 1, 1, "", "WeibullLogLoss"], [130, 4, 1, "", "weibull_log_probs"]], "cyto_dl.nn.losses.weibull.WeibullLogLoss": [[130, 2, 1, "", "forward"]], "cyto_dl.nn.losses.weighted_mse_loss": [[131, 1, 1, "", "WeightedMSELoss"]], "cyto_dl.nn.losses.weighted_mse_loss.WeightedMSELoss": [[131, 2, 1, "", "forward"]], "cyto_dl.nn.mlp": [[132, 1, 1, "", "MLP"]], "cyto_dl.nn.mlp.MLP": [[132, 2, 1, "", "forward"]], "cyto_dl.nn.res_unit": [[138, 1, 1, "", "ResidualUnit"]], "cyto_dl.nn.res_unit.ResidualUnit": [[138, 2, 1, "", "forward"]], "cyto_dl.nn.spatial_transformer": [[139, 1, 1, "", "ConvPoolReLU"], [139, 1, 1, "", "STN"]], "cyto_dl.nn.spatial_transformer.ConvPoolReLU": [[139, 2, 1, "", "forward"]], "cyto_dl.nn.spatial_transformer.STN": [[139, 2, 1, "", "forward"], [139, 2, 1, "", "toggle"]], "cyto_dl.nn.track_sequence_predictor": [[140, 1, 1, "", "TrackClassifier"], [140, 4, 1, "", "positionalencoding1d"]], "cyto_dl.nn.track_sequence_predictor.TrackClassifier": [[140, 2, 1, "", "forward"]], "cyto_dl.nn.vits": [[142, 0, 0, "-", "blocks"], [150, 0, 0, "-", "decoder"], [151, 0, 0, "-", "encoder"], [152, 0, 0, "-", "mae"], [153, 0, 0, "-", "predictor"], [154, 0, 0, "-", "seg"], [155, 0, 0, "-", "utils"]], "cyto_dl.nn.vits.blocks": [[143, 0, 0, "-", "cross_attention"], [144, 0, 0, "-", "intermediate_weigher"], [145, 0, 0, "-", "masked_unit_attention"], [146, 0, 0, "-", "patchify"]], "cyto_dl.nn.vits.blocks.cross_attention": [[143, 1, 1, "", "CrossAttention"], [143, 1, 1, "", "CrossAttentionBlock"], [143, 1, 1, "", "CrossSelfBlock"], [143, 1, 1, "", "Mlp"]], "cyto_dl.nn.vits.blocks.cross_attention.CrossAttention": [[143, 2, 1, "", "forward"]], "cyto_dl.nn.vits.blocks.cross_attention.CrossAttentionBlock": [[143, 2, 1, "", "forward"]], "cyto_dl.nn.vits.blocks.cross_attention.CrossSelfBlock": [[143, 2, 1, "", "forward"]], "cyto_dl.nn.vits.blocks.cross_attention.Mlp": [[143, 2, 1, "", "forward"]], "cyto_dl.nn.vits.blocks.intermediate_weigher": [[144, 1, 1, "", "IntermediateWeigher"]], "cyto_dl.nn.vits.blocks.intermediate_weigher.IntermediateWeigher": [[144, 2, 1, "", "forward"]], "cyto_dl.nn.vits.blocks.masked_unit_attention": [[145, 1, 1, "", "HieraBlock"], [145, 1, 1, "", "MaskUnitAttention"]], "cyto_dl.nn.vits.blocks.masked_unit_attention.HieraBlock": [[145, 2, 1, "", "forward"]], "cyto_dl.nn.vits.blocks.masked_unit_attention.MaskUnitAttention": [[145, 2, 1, "", "forward"]], "cyto_dl.nn.vits.blocks.patchify": [[147, 0, 0, "-", "patchify"], [148, 0, 0, "-", "patchify_base"], [149, 0, 0, "-", "patchify_hiera"]], "cyto_dl.nn.vits.blocks.patchify.patchify": [[147, 1, 1, "", "Patchify"]], "cyto_dl.nn.vits.blocks.patchify.patchify.Patchify": [[147, 2, 1, "", "create_img2token"], [147, 2, 1, "", "extract_visible_tokens"], [147, 2, 1, "", "get_mask_args"], [147, 5, 1, "", "img2token"]], "cyto_dl.nn.vits.blocks.patchify.patchify_base": [[148, 1, 1, "", "PatchifyBase"]], "cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase": [[148, 2, 1, "", "create_conv"], [148, 2, 1, "", "create_patch2img"], [148, 2, 1, "", "extract_visible_tokens"], [148, 2, 1, "", "forward"], [148, 2, 1, "", "get_mask"], [148, 2, 1, "", "get_mask_args"], [148, 5, 1, "", "img2token"]], "cyto_dl.nn.vits.blocks.patchify.patchify_hiera": [[149, 1, 1, "", "PatchifyHiera"], [149, 4, 1, "", "take_indexes_mask"]], "cyto_dl.nn.vits.blocks.patchify.patchify_hiera.PatchifyHiera": [[149, 2, 1, "", "create_img2token"], [149, 2, 1, "", "extract_visible_tokens"], [149, 2, 1, "", "get_mask_args"], [149, 5, 1, "", "img2token"]], "cyto_dl.nn.vits.decoder": [[150, 1, 1, "", "CrossMAE_Decoder"], [150, 1, 1, "", "MAE_Decoder"]], "cyto_dl.nn.vits.decoder.CrossMAE_Decoder": [[150, 2, 1, "", "forward"]], "cyto_dl.nn.vits.decoder.MAE_Decoder": [[150, 2, 1, "", "add_mask_tokens"], [150, 2, 1, "", "adjust_indices_for_cls"], [150, 2, 1, "", "forward"], [150, 2, 1, "", "init_weight"]], "cyto_dl.nn.vits.encoder": [[151, 1, 1, "", "HieraEncoder"], [151, 1, 1, "", "JEPAEncoder"], [151, 1, 1, "", "MAE_Encoder"], [151, 1, 1, "", "SpatialMerger"]], "cyto_dl.nn.vits.encoder.HieraEncoder": [[151, 2, 1, "", "forward"]], "cyto_dl.nn.vits.encoder.JEPAEncoder": [[151, 2, 1, "", "forward"]], "cyto_dl.nn.vits.encoder.MAE_Encoder": [[151, 2, 1, "", "forward"], [151, 2, 1, "", "init_weight"]], "cyto_dl.nn.vits.encoder.SpatialMerger": [[151, 2, 1, "", "forward"]], "cyto_dl.nn.vits.mae": [[152, 1, 1, "", "HieraMAE"], [152, 1, 1, "", "MAE"], [152, 1, 1, "", "MAE_Base"]], "cyto_dl.nn.vits.mae.HieraMAE": [[152, 5, 1, "", "decoder"], [152, 5, 1, "", "encoder"]], "cyto_dl.nn.vits.mae.MAE": [[152, 5, 1, "", "decoder"], [152, 5, 1, "", "encoder"]], "cyto_dl.nn.vits.mae.MAE_Base": [[152, 5, 1, "", "decoder"], [152, 5, 1, "", "encoder"], [152, 2, 1, "", "forward"], [152, 2, 1, "", "init_decoder"], [152, 2, 1, "", "init_encoder"]], "cyto_dl.nn.vits.predictor": [[153, 1, 1, "", "IWMPredictor"], [153, 1, 1, "", "JEPAPredictor"]], "cyto_dl.nn.vits.predictor.IWMPredictor": [[153, 2, 1, "", "forward"]], "cyto_dl.nn.vits.predictor.JEPAPredictor": [[153, 2, 1, "", "forward"], [153, 2, 1, "", "init_weight"], [153, 2, 1, "", "predict_target_features"]], "cyto_dl.nn.vits.seg": [[154, 1, 1, "", "EncodedSkip"], [154, 1, 1, "", "Seg_ViT"], [154, 1, 1, "", "SuperresDecoder"]], "cyto_dl.nn.vits.seg.EncodedSkip": [[154, 2, 1, "", "forward"]], "cyto_dl.nn.vits.seg.Seg_ViT": [[154, 2, 1, "", "forward"]], "cyto_dl.nn.vits.seg.SuperresDecoder": [[154, 2, 1, "", "forward"]], "cyto_dl.nn.vits.utils": [[155, 4, 1, "", "get_positional_embedding"], [155, 4, 1, "", "match_tuple_dimensions"], [155, 4, 1, "", "random_indexes"], [155, 4, 1, "", "take_indexes"]], "cyto_dl.train": [[159, 4, 1, "", "main"]], "cyto_dl.utils": [[161, 0, 0, "-", "array"], [162, 0, 0, "-", "checkpoint"], [163, 0, 0, "-", "config"], [164, 0, 0, "-", "download_test_data"], [165, 0, 0, "-", "dummy_dataset"], [166, 0, 0, "-", "pylogger"], [167, 0, 0, "-", "rich_utils"], [168, 0, 0, "-", "rotation"], [172, 0, 0, "-", "template_utils"]], "cyto_dl.utils.array": [[161, 4, 1, "", "create_dataloader"], [161, 4, 1, "", "extract_array_predictions"]], "cyto_dl.utils.checkpoint": [[162, 4, 1, "", "load_checkpoint"]], "cyto_dl.utils.config": [[163, 4, 1, "", "is_config"], [163, 4, 1, "", "kv_to_dict"], [163, 4, 1, "", "remove_aux_key"]], "cyto_dl.utils.download_test_data": [[164, 4, 1, "", "delete_test_data"], [164, 4, 1, "", "download_test_data"], [164, 4, 1, "", "parse_s3_path"], [164, 4, 1, "", "setup_paths"]], "cyto_dl.utils.dummy_dataset": [[165, 1, 1, "", "DummyDatamodule"], [165, 1, 1, "", "DummyDataset"]], "cyto_dl.utils.dummy_dataset.DummyDatamodule": [[165, 2, 1, "", "get_dataloader"], [165, 2, 1, "", "predict_dataloader"], [165, 2, 1, "", "test_dataloader"], [165, 2, 1, "", "train_dataloader"], [165, 2, 1, "", "val_dataloader"]], "cyto_dl.utils.dummy_dataset.DummyDataset": [[165, 2, 1, "", "generate_img"]], "cyto_dl.utils.pylogger": [[166, 4, 1, "", "get_pylogger"]], "cyto_dl.utils.rich_utils": [[167, 4, 1, "", "enforce_tags"], [167, 4, 1, "", "print_config_tree"]], "cyto_dl.utils.rotation": [[168, 1, 1, "", "RotationModule"], [168, 4, 1, "", "get_rotation_matrix_so2"], [168, 4, 1, "", "get_rotation_matrix_so3"]], "cyto_dl.utils.rotation.RotationModule": [[168, 2, 1, "", "compute_rotation_matrix"]], "cyto_dl.utils.template_utils": [[172, 4, 1, "", "close_loggers"], [172, 4, 1, "", "extras"], [172, 4, 1, "", "get_metric_value"], [172, 4, 1, "", "instantiate_callbacks"], [172, 4, 1, "", "instantiate_loggers"], [172, 4, 1, "", "log_hyperparameters"], [172, 4, 1, "", "save_file"], [172, 4, 1, "", "task_wrapper"]]}, "objtypes": {"0": "py:module", "1": "py:class", "2": "py:method", "3": "py:attribute", "4": "py:function", "5": "py:property"}, "objnames": {"0": ["py", "module", "Python module"], "1": ["py", "class", "Python class"], "2": ["py", "method", "Python method"], "3": ["py", "attribute", "Python attribute"], "4": ["py", "function", "Python function"], "5": ["py", "property", "Python property"]}, "titleterms": {"us": 0, "api": [0, 2, 3, 4, 5, 6, 7], "load": 0, "config": [0, 163, 177], "train": [0, 159, 177], "predict": [0, 177], "util": [0, 30, 75, 76, 77, 78, 79, 80, 81, 86, 87, 88, 155, 160, 161, 162, 163, 164, 165, 166, 167, 168, 169, 170, 171, 172], "cyto_dl": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153, 154, 155, 156, 157, 158, 159, 160, 161, 162, 163, 164, 165, 166, 167, 168, 169, 170, 171, 172, 175], "packag": [1, 2, 3, 8, 17, 19, 24, 27, 36, 37, 44, 58, 60, 63, 66, 68, 72, 75, 78, 82, 86, 89, 95, 99, 105, 106, 109, 118, 133, 141, 142, 146, 156, 157, 160, 169], "subpackag": [1, 2, 17, 24, 36, 60, 72, 75, 89, 105, 141, 142, 156, 160], "submodul": [1, 2, 3, 8, 17, 19, 24, 27, 37, 44, 58, 60, 63, 66, 68, 72, 75, 78, 82, 86, 89, 95, 99, 105, 106, 109, 118, 133, 141, 142, 146, 157, 160, 169], "cyto_dl_model": [3, 4, 5], "cyto_dl_base_model": 4, "modul": [4, 5, 6, 7, 9, 10, 11, 12, 13, 14, 15, 16, 18, 20, 21, 22, 23, 25, 26, 28, 29, 30, 31, 32, 33, 34, 35, 38, 39, 40, 41, 42, 43, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 59, 61, 62, 64, 65, 67, 69, 70, 71, 73, 74, 76, 77, 79, 80, 81, 83, 84, 85, 87, 88, 90, 91, 92, 93, 94, 96, 97, 98, 100, 101, 102, 103, 104, 107, 108, 110, 111, 112, 113, 114, 115, 116, 117, 119, 120, 121, 122, 123, 124, 125, 126, 127, 128, 129, 130, 131, 132, 134, 135, 136, 137, 138, 139, 140, 143, 144, 145, 147, 148, 149, 150, 151, 152, 153, 154, 155, 158, 159, 161, 162, 163, 164, 165, 166, 167, 168, 170, 171, 172], "segmentation_plugin_model": 5, "data": [6, 173, 176], "model": [7, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 173, 176], "callback": [8, 9, 10, 11, 12, 13, 14, 15], "callback_util": 9, "csv_saver": 10, "image_sav": 11, "latent_walk": 12, "layer_freez": 13, "model_util": 14, "outlier_detect": 15, "compil": 16, "datafram": [17, 18, 19, 20, 21, 22, 23, 27, 28, 29, 30], "reader": 18, "transform": [19, 20, 21, 22, 23, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57], "filter": 20, "group_col": 21, "misc": 22, "split": 23, "datamodul": [24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34], "arrai": [25, 161], "data_dict": 26, "dataframe_datamodul": 28, "grouped_dataframe_datamodul": 29, "folder": 31, "multidim_imag": 32, "smartcach": 33, "torchvis": 34, "eval": 35, "imag": [36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57], "io": [37, 38, 39, 40, 41, 42, 43, 157, 158], "aicsimage_load": 38, "monai_bio_read": 39, "numpy_read": 40, "ome_zarr_read": 41, "polygon_load": 42, "skimage_read": 43, "add_meta": 45, "bright_sampl": 46, "clip": 47, "contrastadjust": 48, "generate_jepa_mask": 49, "merg": 50, "multiscale_cropp": 51, "pad": 52, "project": 53, "rotation_mask_transform": 54, "save": 55, "so2_random_rot": 56, "track_transform": 57, "logger": [58, 59], "mlflow": [59, 87], "base_model": 61, "basic_model": 62, "classif": [63, 64, 65], "timepoint_classif": 65, "contrast": [66, 67], "handler": [68, 69, 70, 71], "base_handl": 69, "image_handl": 70, "load_image_patch": 71, "im2im": [72, 73, 74, 75, 76, 77, 78, 79, 80, 81], "gan": 73, "multi_task": 74, "instance_seg": 76, "noise_anneal": 77, "postprocess": [78, 79, 80, 81], "act_thresh_label": 79, "auto_thresh": 80, "dict_to_im": 81, "jepa": [82, 83, 84, 85], "ijepa": 83, "iwm": 84, "jepa_bas": 85, "vae": [89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104], "base_va": 90, "image_canon_va": 91, "image_encod": 92, "image_va": 93, "latent_loss_va": 94, "o2_spharm_va": [95, 96, 97], "o2_spharm_encod": 96, "point_cloud_va": 98, "prior": [99, 100, 101, 102, 103], "abstract_prior": 100, "gaussian": 101, "identity_prior": 102, "joint_prior": 103, "tabular_va": 104, "nn": [105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153, 154, 155], "discrimin": [106, 107, 108], "multi_scale_discrimin": 107, "n_layer_discrimin": 108, "head": [109, 110, 111, 112, 113, 114, 115, 116], "base_head": 110, "gan_head": 111, "gan_head_superr": 112, "mae_head": 113, "mask_head": 114, "res_blocks_head": 115, "vic_reg": [116, 129], "hr_skip": 117, "loss": [118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 128, 129, 130, 131], "adversarial_loss": 119, "chamfer_loss": 120, "continuous_bernoulli": 121, "cosine_loss": 122, "gan_loss": 123, "gaussian_nll_loss": 124, "geomloss": 125, "loss_wrapp": 126, "spharm_loss": 127, "threshold_loss": 128, "weibul": 130, "weighted_mse_loss": 131, "mlp": 132, "point_cloud": [133, 134, 135, 136, 137, 156, 157, 158], "dgcnn": 134, "folding_net": 135, "graph_funct": 136, "vnn": 137, "res_unit": 138, "spatial_transform": 139, "track_sequence_predictor": 140, "vit": [141, 142, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153, 154, 155], "block": [142, 143, 144, 145, 146, 147, 148, 149], "cross_attent": 143, "intermediate_weigh": 144, "masked_unit_attent": 145, "patchifi": [146, 147, 148, 149], "patchify_bas": 148, "patchify_hiera": 149, "decod": 150, "encod": 151, "mae": 152, "predictor": 153, "seg": 154, "read_pcloud": 158, "checkpoint": 162, "download_test_data": 164, "dummy_dataset": 165, "pylogg": 166, "rich_util": 167, "rotat": [168, 171], "spharm": [169, 170, 171], "reconstruct": 170, "template_util": 172, "overview": [173, 176], "experi": [173, 176], "instal": 174, "modifi": 177, "exampl": 177, "test": 177}, "envversion": {"sphinx.domains.c": 2, "sphinx.domains.changeset": 1, "sphinx.domains.citation": 1, "sphinx.domains.cpp": 8, "sphinx.domains.index": 1, "sphinx.domains.javascript": 2, "sphinx.domains.math": 2, "sphinx.domains.python": 3, "sphinx.domains.rst": 2, "sphinx.domains.std": 2, "sphinx.ext.viewcode": 1, "sphinx": 57}, "alltitles": {"Using the API": [[0, "using-the-api"]], "Loading Configs": [[0, "loading-configs"]], "Training and Prediction": [[0, "training-and-prediction"]], "Utils": [[0, "utils"]], "cyto_dl package": [[1, "module-cyto_dl"]], "Subpackages": [[1, "subpackages"], [2, "subpackages"], [17, "subpackages"], [24, "subpackages"], [36, "subpackages"], [60, "subpackages"], [72, "subpackages"], [75, "subpackages"], [89, "subpackages"], [105, "subpackages"], [141, "subpackages"], [142, "subpackages"], [156, "subpackages"], [160, "subpackages"]], "Submodules": [[1, "submodules"], [2, "submodules"], [3, "submodules"], [8, "submodules"], [17, "submodules"], [19, "submodules"], [24, "submodules"], [27, "submodules"], [37, "submodules"], [44, "submodules"], [58, "submodules"], [60, "submodules"], [63, "submodules"], [66, "submodules"], [68, "submodules"], [72, "submodules"], [75, "submodules"], [78, "submodules"], [82, "submodules"], [86, "submodules"], [89, "submodules"], [95, "submodules"], [99, "submodules"], [105, "submodules"], [106, "submodules"], [109, "submodules"], [118, "submodules"], [133, "submodules"], [141, "submodules"], [142, "submodules"], [146, "submodules"], [157, "submodules"], [160, "submodules"], [169, "submodules"]], "cyto_dl.api package": [[2, "module-cyto_dl.api"]], "cyto_dl.api.cyto_dl_model package": [[3, "module-cyto_dl.api.cyto_dl_model"]], "cyto_dl.api.cyto_dl_model.cyto_dl_base_model module": [[4, "module-cyto_dl.api.cyto_dl_model.cyto_dl_base_model"]], "cyto_dl.api.cyto_dl_model.segmentation_plugin_model module": [[5, "module-cyto_dl.api.cyto_dl_model.segmentation_plugin_model"]], "cyto_dl.api.data module": [[6, "module-cyto_dl.api.data"]], "cyto_dl.api.model module": [[7, "module-cyto_dl.api.model"]], "cyto_dl.callbacks package": [[8, "module-cyto_dl.callbacks"]], "cyto_dl.callbacks.callback_utils module": [[9, "cyto-dl-callbacks-callback-utils-module"]], "cyto_dl.callbacks.csv_saver module": [[10, "module-cyto_dl.callbacks.csv_saver"]], "cyto_dl.callbacks.image_saver module": [[11, "module-cyto_dl.callbacks.image_saver"]], "cyto_dl.callbacks.latent_walk module": [[12, "cyto-dl-callbacks-latent-walk-module"]], "cyto_dl.callbacks.layer_freeze module": [[13, "module-cyto_dl.callbacks.layer_freeze"]], "cyto_dl.callbacks.model_utils module": [[14, "module-cyto_dl.callbacks.model_utils"]], "cyto_dl.callbacks.outlier_detection module": [[15, "module-cyto_dl.callbacks.outlier_detection"]], "cyto_dl.compile module": [[16, "cyto-dl-compile-module"]], "cyto_dl.dataframe package": [[17, "module-cyto_dl.dataframe"]], "cyto_dl.dataframe.readers module": [[18, "module-cyto_dl.dataframe.readers"]], "cyto_dl.dataframe.transforms package": [[19, "module-cyto_dl.dataframe.transforms"]], "cyto_dl.dataframe.transforms.filter module": [[20, "module-cyto_dl.dataframe.transforms.filter"]], "cyto_dl.dataframe.transforms.group_cols module": [[21, "module-cyto_dl.dataframe.transforms.group_cols"]], "cyto_dl.dataframe.transforms.misc module": [[22, "module-cyto_dl.dataframe.transforms.misc"]], "cyto_dl.dataframe.transforms.split module": [[23, "module-cyto_dl.dataframe.transforms.split"]], "cyto_dl.datamodules package": [[24, "module-cyto_dl.datamodules"]], "cyto_dl.datamodules.array module": [[25, "module-cyto_dl.datamodules.array"]], "cyto_dl.datamodules.data_dict module": [[26, "module-cyto_dl.datamodules.data_dict"]], "cyto_dl.datamodules.dataframe package": [[27, "module-cyto_dl.datamodules.dataframe"]], "cyto_dl.datamodules.dataframe.dataframe_datamodule module": [[28, "module-cyto_dl.datamodules.dataframe.dataframe_datamodule"]], "cyto_dl.datamodules.dataframe.grouped_dataframe_datamodule module": [[29, "module-cyto_dl.datamodules.dataframe.grouped_dataframe_datamodule"]], "cyto_dl.datamodules.dataframe.utils module": [[30, "module-cyto_dl.datamodules.dataframe.utils"]], "cyto_dl.datamodules.folder module": [[31, "module-cyto_dl.datamodules.folder"]], "cyto_dl.datamodules.multidim_image module": [[32, "module-cyto_dl.datamodules.multidim_image"]], "cyto_dl.datamodules.smartcache module": [[33, "module-cyto_dl.datamodules.smartcache"]], "cyto_dl.datamodules.torchvision module": [[34, "module-cyto_dl.datamodules.torchvision"]], "cyto_dl.eval module": [[35, "module-cyto_dl.eval"]], "cyto_dl.image package": [[36, "module-cyto_dl.image"]], "cyto_dl.image.io package": [[37, "module-cyto_dl.image.io"]], "cyto_dl.image.io.aicsimage_loader module": [[38, "module-cyto_dl.image.io.aicsimage_loader"]], "cyto_dl.image.io.monai_bio_reader module": [[39, "module-cyto_dl.image.io.monai_bio_reader"]], "cyto_dl.image.io.numpy_reader module": [[40, "module-cyto_dl.image.io.numpy_reader"]], "cyto_dl.image.io.ome_zarr_reader module": [[41, "module-cyto_dl.image.io.ome_zarr_reader"]], "cyto_dl.image.io.polygon_loader module": [[42, "module-cyto_dl.image.io.polygon_loader"]], "cyto_dl.image.io.skimage_reader module": [[43, "module-cyto_dl.image.io.skimage_reader"]], "cyto_dl.image.transforms package": [[44, "module-cyto_dl.image.transforms"]], "cyto_dl.image.transforms.add_meta module": [[45, "module-cyto_dl.image.transforms.add_meta"]], "cyto_dl.image.transforms.bright_sampler module": [[46, "module-cyto_dl.image.transforms.bright_sampler"]], "cyto_dl.image.transforms.clip module": [[47, "module-cyto_dl.image.transforms.clip"]], "cyto_dl.image.transforms.contrastadjust module": [[48, "module-cyto_dl.image.transforms.contrastadjust"]], "cyto_dl.image.transforms.generate_jepa_masks module": [[49, "module-cyto_dl.image.transforms.generate_jepa_masks"]], "cyto_dl.image.transforms.merge module": [[50, "module-cyto_dl.image.transforms.merge"]], "cyto_dl.image.transforms.multiscale_cropper module": [[51, "module-cyto_dl.image.transforms.multiscale_cropper"]], "cyto_dl.image.transforms.pad module": [[52, "module-cyto_dl.image.transforms.pad"]], "cyto_dl.image.transforms.project module": [[53, "module-cyto_dl.image.transforms.project"]], "cyto_dl.image.transforms.rotation_mask_transform module": [[54, "cyto-dl-image-transforms-rotation-mask-transform-module"]], "cyto_dl.image.transforms.save module": [[55, "module-cyto_dl.image.transforms.save"]], "cyto_dl.image.transforms.so2_random_rotation module": [[56, "cyto-dl-image-transforms-so2-random-rotation-module"]], "cyto_dl.image.transforms.track_transforms module": [[57, "module-cyto_dl.image.transforms.track_transforms"]], "cyto_dl.loggers package": [[58, "module-cyto_dl.loggers"]], "cyto_dl.loggers.mlflow module": [[59, "module-cyto_dl.loggers.mlflow"]], "cyto_dl.models package": [[60, "module-cyto_dl.models"]], "cyto_dl.models.base_model module": [[61, "module-cyto_dl.models.base_model"]], "cyto_dl.models.basic_model module": [[62, "module-cyto_dl.models.basic_model"]], "cyto_dl.models.classification package": [[63, "module-cyto_dl.models.classification"]], "cyto_dl.models.classification.classification module": [[64, "module-cyto_dl.models.classification.classification"]], "cyto_dl.models.classification.timepoint_classification module": [[65, "module-cyto_dl.models.classification.timepoint_classification"]], "cyto_dl.models.contrastive package": [[66, "module-cyto_dl.models.contrastive"]], "cyto_dl.models.contrastive.contrastive module": [[67, "module-cyto_dl.models.contrastive.contrastive"]], "cyto_dl.models.handlers package": [[68, "cyto-dl-models-handlers-package"]], "cyto_dl.models.handlers.base_handler module": [[69, "cyto-dl-models-handlers-base-handler-module"]], "cyto_dl.models.handlers.image_handler module": [[70, "cyto-dl-models-handlers-image-handler-module"]], "cyto_dl.models.handlers.load_image_patch module": [[71, "cyto-dl-models-handlers-load-image-patch-module"]], "cyto_dl.models.im2im package": [[72, "module-cyto_dl.models.im2im"]], "cyto_dl.models.im2im.gan module": [[73, "module-cyto_dl.models.im2im.gan"]], "cyto_dl.models.im2im.multi_task module": [[74, "module-cyto_dl.models.im2im.multi_task"]], "cyto_dl.models.im2im.utils package": [[75, "module-cyto_dl.models.im2im.utils"]], "cyto_dl.models.im2im.utils.instance_seg module": [[76, "module-cyto_dl.models.im2im.utils.instance_seg"]], "cyto_dl.models.im2im.utils.noise_annealer module": [[77, "module-cyto_dl.models.im2im.utils.noise_annealer"]], "cyto_dl.models.im2im.utils.postprocessing package": [[78, "module-cyto_dl.models.im2im.utils.postprocessing"]], "cyto_dl.models.im2im.utils.postprocessing.act_thresh_label module": [[79, "module-cyto_dl.models.im2im.utils.postprocessing.act_thresh_label"]], "cyto_dl.models.im2im.utils.postprocessing.auto_thresh module": [[80, "module-cyto_dl.models.im2im.utils.postprocessing.auto_thresh"]], "cyto_dl.models.im2im.utils.postprocessing.dict_to_im module": [[81, "module-cyto_dl.models.im2im.utils.postprocessing.dict_to_im"]], "cyto_dl.models.jepa package": [[82, "module-cyto_dl.models.jepa"]], "cyto_dl.models.jepa.ijepa module": [[83, "module-cyto_dl.models.jepa.ijepa"]], "cyto_dl.models.jepa.iwm module": [[84, "module-cyto_dl.models.jepa.iwm"]], "cyto_dl.models.jepa.jepa_base module": [[85, "module-cyto_dl.models.jepa.jepa_base"]], "cyto_dl.models.utils package": [[86, "module-cyto_dl.models.utils"]], "cyto_dl.models.utils.mlflow module": [[87, "module-cyto_dl.models.utils.mlflow"]], "cyto_dl.models.utils.utils module": [[88, "module-cyto_dl.models.utils.utils"]], "cyto_dl.models.vae package": [[89, "cyto-dl-models-vae-package"]], "cyto_dl.models.vae.base_vae module": [[90, "module-cyto_dl.models.vae.base_vae"]], "cyto_dl.models.vae.image_canon_vae module": [[91, "cyto-dl-models-vae-image-canon-vae-module"]], "cyto_dl.models.vae.image_encoder module": [[92, "cyto-dl-models-vae-image-encoder-module"]], "cyto_dl.models.vae.image_vae module": [[93, "cyto-dl-models-vae-image-vae-module"]], "cyto_dl.models.vae.latent_loss_vae module": [[94, "cyto-dl-models-vae-latent-loss-vae-module"]], "cyto_dl.models.vae.o2_spharm_vae package": [[95, "cyto-dl-models-vae-o2-spharm-vae-package"]], "cyto_dl.models.vae.o2_spharm_vae.o2_spharm_encoder module": [[96, "cyto-dl-models-vae-o2-spharm-vae-o2-spharm-encoder-module"]], "cyto_dl.models.vae.o2_spharm_vae.o2_spharm_vae module": [[97, "cyto-dl-models-vae-o2-spharm-vae-o2-spharm-vae-module"]], "cyto_dl.models.vae.point_cloud_vae module": [[98, "cyto-dl-models-vae-point-cloud-vae-module"]], "cyto_dl.models.vae.priors package": [[99, "module-cyto_dl.models.vae.priors"]], "cyto_dl.models.vae.priors.abstract_prior module": [[100, "module-cyto_dl.models.vae.priors.abstract_prior"]], "cyto_dl.models.vae.priors.gaussian module": [[101, "module-cyto_dl.models.vae.priors.gaussian"]], "cyto_dl.models.vae.priors.identity_prior module": [[102, "module-cyto_dl.models.vae.priors.identity_prior"]], "cyto_dl.models.vae.priors.joint_prior module": [[103, "module-cyto_dl.models.vae.priors.joint_prior"]], "cyto_dl.models.vae.tabular_vae module": [[104, "cyto-dl-models-vae-tabular-vae-module"]], "cyto_dl.nn package": [[105, "module-cyto_dl.nn"]], "cyto_dl.nn.discriminators package": [[106, "module-cyto_dl.nn.discriminators"]], "cyto_dl.nn.discriminators.multi_scale_discriminator module": [[107, "module-cyto_dl.nn.discriminators.multi_scale_discriminator"]], "cyto_dl.nn.discriminators.n_layer_discriminator module": [[108, "module-cyto_dl.nn.discriminators.n_layer_discriminator"]], "cyto_dl.nn.head package": [[109, "module-cyto_dl.nn.head"]], "cyto_dl.nn.head.base_head module": [[110, "module-cyto_dl.nn.head.base_head"]], "cyto_dl.nn.head.gan_head module": [[111, "module-cyto_dl.nn.head.gan_head"]], "cyto_dl.nn.head.gan_head_superres module": [[112, "module-cyto_dl.nn.head.gan_head_superres"]], "cyto_dl.nn.head.mae_head module": [[113, "module-cyto_dl.nn.head.mae_head"]], "cyto_dl.nn.head.mask_head module": [[114, "module-cyto_dl.nn.head.mask_head"]], "cyto_dl.nn.head.res_blocks_head module": [[115, "module-cyto_dl.nn.head.res_blocks_head"]], "cyto_dl.nn.head.vic_reg module": [[116, "module-cyto_dl.nn.head.vic_reg"]], "cyto_dl.nn.hr_skip module": [[117, "module-cyto_dl.nn.hr_skip"]], "cyto_dl.nn.losses package": [[118, "module-cyto_dl.nn.losses"]], "cyto_dl.nn.losses.adversarial_loss module": [[119, "module-cyto_dl.nn.losses.adversarial_loss"]], "cyto_dl.nn.losses.chamfer_loss module": [[120, "module-cyto_dl.nn.losses.chamfer_loss"]], "cyto_dl.nn.losses.continuous_bernoulli module": [[121, "module-cyto_dl.nn.losses.continuous_bernoulli"]], "cyto_dl.nn.losses.cosine_loss module": [[122, "module-cyto_dl.nn.losses.cosine_loss"]], "cyto_dl.nn.losses.gan_loss module": [[123, "module-cyto_dl.nn.losses.gan_loss"]], "cyto_dl.nn.losses.gaussian_nll_loss module": [[124, "module-cyto_dl.nn.losses.gaussian_nll_loss"]], "cyto_dl.nn.losses.geomloss module": [[125, "cyto-dl-nn-losses-geomloss-module"]], "cyto_dl.nn.losses.loss_wrapper module": [[126, "module-cyto_dl.nn.losses.loss_wrapper"]], "cyto_dl.nn.losses.spharm_loss module": [[127, "cyto-dl-nn-losses-spharm-loss-module"]], "cyto_dl.nn.losses.threshold_loss module": [[128, "module-cyto_dl.nn.losses.threshold_loss"]], "cyto_dl.nn.losses.vic_reg module": [[129, "module-cyto_dl.nn.losses.vic_reg"]], "cyto_dl.nn.losses.weibull module": [[130, "module-cyto_dl.nn.losses.weibull"]], "cyto_dl.nn.losses.weighted_mse_loss module": [[131, "module-cyto_dl.nn.losses.weighted_mse_loss"]], "cyto_dl.nn.mlp module": [[132, "module-cyto_dl.nn.mlp"]], "cyto_dl.nn.point_cloud package": [[133, "cyto-dl-nn-point-cloud-package"]], "cyto_dl.nn.point_cloud.dgcnn module": [[134, "cyto-dl-nn-point-cloud-dgcnn-module"]], "cyto_dl.nn.point_cloud.folding_net module": [[135, "cyto-dl-nn-point-cloud-folding-net-module"]], "cyto_dl.nn.point_cloud.graph_functions module": [[136, "cyto-dl-nn-point-cloud-graph-functions-module"]], "cyto_dl.nn.point_cloud.vnn module": [[137, "cyto-dl-nn-point-cloud-vnn-module"]], "cyto_dl.nn.res_unit module": [[138, "module-cyto_dl.nn.res_unit"]], "cyto_dl.nn.spatial_transformer module": [[139, "module-cyto_dl.nn.spatial_transformer"]], "cyto_dl.nn.track_sequence_predictor module": [[140, "module-cyto_dl.nn.track_sequence_predictor"]], "cyto_dl.nn.vits package": [[141, "module-cyto_dl.nn.vits"]], "cyto_dl.nn.vits.blocks package": [[142, "module-cyto_dl.nn.vits.blocks"]], "cyto_dl.nn.vits.blocks.cross_attention module": [[143, "module-cyto_dl.nn.vits.blocks.cross_attention"]], "cyto_dl.nn.vits.blocks.intermediate_weigher module": [[144, "module-cyto_dl.nn.vits.blocks.intermediate_weigher"]], "cyto_dl.nn.vits.blocks.masked_unit_attention module": [[145, "module-cyto_dl.nn.vits.blocks.masked_unit_attention"]], "cyto_dl.nn.vits.blocks.patchify package": [[146, "module-cyto_dl.nn.vits.blocks.patchify"]], "cyto_dl.nn.vits.blocks.patchify.patchify module": [[147, "module-cyto_dl.nn.vits.blocks.patchify.patchify"]], "cyto_dl.nn.vits.blocks.patchify.patchify_base module": [[148, "module-cyto_dl.nn.vits.blocks.patchify.patchify_base"]], "cyto_dl.nn.vits.blocks.patchify.patchify_hiera module": [[149, "module-cyto_dl.nn.vits.blocks.patchify.patchify_hiera"]], "cyto_dl.nn.vits.decoder module": [[150, "module-cyto_dl.nn.vits.decoder"]], "cyto_dl.nn.vits.encoder module": [[151, "module-cyto_dl.nn.vits.encoder"]], "cyto_dl.nn.vits.mae module": [[152, "module-cyto_dl.nn.vits.mae"]], "cyto_dl.nn.vits.predictor module": [[153, "module-cyto_dl.nn.vits.predictor"]], "cyto_dl.nn.vits.seg module": [[154, "module-cyto_dl.nn.vits.seg"]], "cyto_dl.nn.vits.utils module": [[155, "module-cyto_dl.nn.vits.utils"]], "cyto_dl.point_cloud package": [[156, "module-cyto_dl.point_cloud"]], "cyto_dl.point_cloud.io package": [[157, "cyto-dl-point-cloud-io-package"]], "cyto_dl.point_cloud.io.read_pcloud module": [[158, "cyto-dl-point-cloud-io-read-pcloud-module"]], "cyto_dl.train module": [[159, "module-cyto_dl.train"]], "cyto_dl.utils package": [[160, "module-cyto_dl.utils"]], "cyto_dl.utils.array module": [[161, "module-cyto_dl.utils.array"]], "cyto_dl.utils.checkpoint module": [[162, "module-cyto_dl.utils.checkpoint"]], "cyto_dl.utils.config module": [[163, "module-cyto_dl.utils.config"]], "cyto_dl.utils.download_test_data module": [[164, "module-cyto_dl.utils.download_test_data"]], "cyto_dl.utils.dummy_dataset module": [[165, "module-cyto_dl.utils.dummy_dataset"]], "cyto_dl.utils.pylogger module": [[166, "module-cyto_dl.utils.pylogger"]], "cyto_dl.utils.rich_utils module": [[167, "module-cyto_dl.utils.rich_utils"]], "cyto_dl.utils.rotation module": [[168, "module-cyto_dl.utils.rotation"]], "cyto_dl.utils.spharm package": [[169, "cyto-dl-utils-spharm-package"]], "cyto_dl.utils.spharm.reconstruction module": [[170, "cyto-dl-utils-spharm-reconstruction-module"]], "cyto_dl.utils.spharm.rotation module": [[171, "cyto-dl-utils-spharm-rotation-module"]], "cyto_dl.utils.template_utils module": [[172, "module-cyto_dl.utils.template_utils"]], "Overview": [[173, "overview"], [176, "overview"]], "data": [[173, "data"], [176, "data"]], "experiment": [[173, "experiment"], [176, "experiment"]], "model": [[173, "model"], [176, "model"]], "Installation": [[174, "installation"]], "cyto_dl": [[175, "cyto-dl"]], "Modifying the Example Configs": [[177, "modifying-the-example-configs"]], "Training": [[177, "training"]], "Testing/Prediction": [[177, "testing-prediction"]]}, "indexentries": {"cyto_dl": [[1, "module-cyto_dl"]], "module": [[1, "module-cyto_dl"], [2, "module-cyto_dl.api"], [3, "module-cyto_dl.api.cyto_dl_model"], [4, "module-cyto_dl.api.cyto_dl_model.cyto_dl_base_model"], [5, "module-cyto_dl.api.cyto_dl_model.segmentation_plugin_model"], [6, "module-cyto_dl.api.data"], [7, "module-cyto_dl.api.model"], [8, "module-cyto_dl.callbacks"], [10, "module-cyto_dl.callbacks.csv_saver"], [11, "module-cyto_dl.callbacks.image_saver"], [13, "module-cyto_dl.callbacks.layer_freeze"], [14, "module-cyto_dl.callbacks.model_utils"], [15, "module-cyto_dl.callbacks.outlier_detection"], [17, "module-cyto_dl.dataframe"], [18, "module-cyto_dl.dataframe.readers"], [19, "module-cyto_dl.dataframe.transforms"], [20, "module-cyto_dl.dataframe.transforms.filter"], [21, "module-cyto_dl.dataframe.transforms.group_cols"], [22, "module-cyto_dl.dataframe.transforms.misc"], [23, "module-cyto_dl.dataframe.transforms.split"], [24, "module-cyto_dl.datamodules"], [25, "module-cyto_dl.datamodules.array"], [26, "module-cyto_dl.datamodules.data_dict"], [27, "module-cyto_dl.datamodules.dataframe"], [28, "module-cyto_dl.datamodules.dataframe.dataframe_datamodule"], [29, "module-cyto_dl.datamodules.dataframe.grouped_dataframe_datamodule"], [30, "module-cyto_dl.datamodules.dataframe.utils"], [31, "module-cyto_dl.datamodules.folder"], [32, "module-cyto_dl.datamodules.multidim_image"], [33, "module-cyto_dl.datamodules.smartcache"], [34, "module-cyto_dl.datamodules.torchvision"], [35, "module-cyto_dl.eval"], [36, "module-cyto_dl.image"], [37, "module-cyto_dl.image.io"], [38, "module-cyto_dl.image.io.aicsimage_loader"], [39, "module-cyto_dl.image.io.monai_bio_reader"], [40, "module-cyto_dl.image.io.numpy_reader"], [41, "module-cyto_dl.image.io.ome_zarr_reader"], [42, "module-cyto_dl.image.io.polygon_loader"], [43, "module-cyto_dl.image.io.skimage_reader"], [44, "module-cyto_dl.image.transforms"], [45, "module-cyto_dl.image.transforms.add_meta"], [46, "module-cyto_dl.image.transforms.bright_sampler"], [47, "module-cyto_dl.image.transforms.clip"], [48, "module-cyto_dl.image.transforms.contrastadjust"], [49, "module-cyto_dl.image.transforms.generate_jepa_masks"], [50, "module-cyto_dl.image.transforms.merge"], [51, "module-cyto_dl.image.transforms.multiscale_cropper"], [52, "module-cyto_dl.image.transforms.pad"], [53, "module-cyto_dl.image.transforms.project"], [55, "module-cyto_dl.image.transforms.save"], [57, "module-cyto_dl.image.transforms.track_transforms"], [58, "module-cyto_dl.loggers"], [59, "module-cyto_dl.loggers.mlflow"], [60, "module-cyto_dl.models"], [61, "module-cyto_dl.models.base_model"], [62, "module-cyto_dl.models.basic_model"], [63, "module-cyto_dl.models.classification"], [64, "module-cyto_dl.models.classification.classification"], [65, "module-cyto_dl.models.classification.timepoint_classification"], [66, "module-cyto_dl.models.contrastive"], [67, "module-cyto_dl.models.contrastive.contrastive"], [72, "module-cyto_dl.models.im2im"], [73, "module-cyto_dl.models.im2im.gan"], [74, "module-cyto_dl.models.im2im.multi_task"], [75, "module-cyto_dl.models.im2im.utils"], [76, "module-cyto_dl.models.im2im.utils.instance_seg"], [77, "module-cyto_dl.models.im2im.utils.noise_annealer"], [78, "module-cyto_dl.models.im2im.utils.postprocessing"], [79, "module-cyto_dl.models.im2im.utils.postprocessing.act_thresh_label"], [80, "module-cyto_dl.models.im2im.utils.postprocessing.auto_thresh"], [81, "module-cyto_dl.models.im2im.utils.postprocessing.dict_to_im"], [82, "module-cyto_dl.models.jepa"], [83, "module-cyto_dl.models.jepa.ijepa"], [84, "module-cyto_dl.models.jepa.iwm"], [85, "module-cyto_dl.models.jepa.jepa_base"], [86, "module-cyto_dl.models.utils"], [87, "module-cyto_dl.models.utils.mlflow"], [88, "module-cyto_dl.models.utils.utils"], [90, "module-cyto_dl.models.vae.base_vae"], [99, "module-cyto_dl.models.vae.priors"], [100, "module-cyto_dl.models.vae.priors.abstract_prior"], [101, "module-cyto_dl.models.vae.priors.gaussian"], [102, "module-cyto_dl.models.vae.priors.identity_prior"], [103, "module-cyto_dl.models.vae.priors.joint_prior"], [105, "module-cyto_dl.nn"], [106, "module-cyto_dl.nn.discriminators"], [107, "module-cyto_dl.nn.discriminators.multi_scale_discriminator"], [108, "module-cyto_dl.nn.discriminators.n_layer_discriminator"], [109, "module-cyto_dl.nn.head"], [110, "module-cyto_dl.nn.head.base_head"], [111, "module-cyto_dl.nn.head.gan_head"], [112, "module-cyto_dl.nn.head.gan_head_superres"], [113, "module-cyto_dl.nn.head.mae_head"], [114, "module-cyto_dl.nn.head.mask_head"], [115, "module-cyto_dl.nn.head.res_blocks_head"], [116, "module-cyto_dl.nn.head.vic_reg"], [117, "module-cyto_dl.nn.hr_skip"], [118, "module-cyto_dl.nn.losses"], [119, "module-cyto_dl.nn.losses.adversarial_loss"], [120, "module-cyto_dl.nn.losses.chamfer_loss"], [121, "module-cyto_dl.nn.losses.continuous_bernoulli"], [122, "module-cyto_dl.nn.losses.cosine_loss"], [123, "module-cyto_dl.nn.losses.gan_loss"], [124, "module-cyto_dl.nn.losses.gaussian_nll_loss"], [126, "module-cyto_dl.nn.losses.loss_wrapper"], [128, "module-cyto_dl.nn.losses.threshold_loss"], [129, "module-cyto_dl.nn.losses.vic_reg"], [130, "module-cyto_dl.nn.losses.weibull"], [131, "module-cyto_dl.nn.losses.weighted_mse_loss"], [132, "module-cyto_dl.nn.mlp"], [138, "module-cyto_dl.nn.res_unit"], [139, "module-cyto_dl.nn.spatial_transformer"], [140, "module-cyto_dl.nn.track_sequence_predictor"], [141, "module-cyto_dl.nn.vits"], [142, "module-cyto_dl.nn.vits.blocks"], [143, "module-cyto_dl.nn.vits.blocks.cross_attention"], [144, "module-cyto_dl.nn.vits.blocks.intermediate_weigher"], [145, "module-cyto_dl.nn.vits.blocks.masked_unit_attention"], [146, "module-cyto_dl.nn.vits.blocks.patchify"], [147, "module-cyto_dl.nn.vits.blocks.patchify.patchify"], [148, "module-cyto_dl.nn.vits.blocks.patchify.patchify_base"], [149, "module-cyto_dl.nn.vits.blocks.patchify.patchify_hiera"], [150, "module-cyto_dl.nn.vits.decoder"], [151, "module-cyto_dl.nn.vits.encoder"], [152, "module-cyto_dl.nn.vits.mae"], [153, "module-cyto_dl.nn.vits.predictor"], [154, "module-cyto_dl.nn.vits.seg"], [155, "module-cyto_dl.nn.vits.utils"], [156, "module-cyto_dl.point_cloud"], [159, "module-cyto_dl.train"], [160, "module-cyto_dl.utils"], [161, "module-cyto_dl.utils.array"], [162, "module-cyto_dl.utils.checkpoint"], [163, "module-cyto_dl.utils.config"], [164, "module-cyto_dl.utils.download_test_data"], [165, "module-cyto_dl.utils.dummy_dataset"], [166, "module-cyto_dl.utils.pylogger"], [167, "module-cyto_dl.utils.rich_utils"], [168, "module-cyto_dl.utils.rotation"], [172, "module-cyto_dl.utils.template_utils"]], "cyto_dl.api": [[2, "module-cyto_dl.api"]], "cyto_dl.api.cyto_dl_model": [[3, "module-cyto_dl.api.cyto_dl_model"]], "cytodlbasemodel (class in cyto_dl.api.cyto_dl_model.cyto_dl_base_model)": [[4, "cyto_dl.api.cyto_dl_model.cyto_dl_base_model.CytoDLBaseModel"]], "cyto_dl.api.cyto_dl_model.cyto_dl_base_model": [[4, "module-cyto_dl.api.cyto_dl_model.cyto_dl_base_model"]], "from_default_config() (cyto_dl.api.cyto_dl_model.cyto_dl_base_model.cytodlbasemodel class method)": [[4, "cyto_dl.api.cyto_dl_model.cyto_dl_base_model.CytoDLBaseModel.from_default_config"]], "from_existing_config() (cyto_dl.api.cyto_dl_model.cyto_dl_base_model.cytodlbasemodel class method)": [[4, "cyto_dl.api.cyto_dl_model.cyto_dl_base_model.CytoDLBaseModel.from_existing_config"]], "get_config() (cyto_dl.api.cyto_dl_model.cyto_dl_base_model.cytodlbasemodel method)": [[4, "cyto_dl.api.cyto_dl_model.cyto_dl_base_model.CytoDLBaseModel.get_config"]], "get_experiment_name() (cyto_dl.api.cyto_dl_model.cyto_dl_base_model.cytodlbasemodel method)": [[4, "cyto_dl.api.cyto_dl_model.cyto_dl_base_model.CytoDLBaseModel.get_experiment_name"]], "predict() (cyto_dl.api.cyto_dl_model.cyto_dl_base_model.cytodlbasemodel method)": [[4, "cyto_dl.api.cyto_dl_model.cyto_dl_base_model.CytoDLBaseModel.predict"]], "save_config() (cyto_dl.api.cyto_dl_model.cyto_dl_base_model.cytodlbasemodel method)": [[4, "cyto_dl.api.cyto_dl_model.cyto_dl_base_model.CytoDLBaseModel.save_config"]], "set_experiment_name() (cyto_dl.api.cyto_dl_model.cyto_dl_base_model.cytodlbasemodel method)": [[4, "cyto_dl.api.cyto_dl_model.cyto_dl_base_model.CytoDLBaseModel.set_experiment_name"]], "train() (cyto_dl.api.cyto_dl_model.cyto_dl_base_model.cytodlbasemodel method)": [[4, "cyto_dl.api.cyto_dl_model.cyto_dl_base_model.CytoDLBaseModel.train"]], "segmentationpluginmodel (class in cyto_dl.api.cyto_dl_model.segmentation_plugin_model)": [[5, "cyto_dl.api.cyto_dl_model.segmentation_plugin_model.SegmentationPluginModel"]], "cyto_dl.api.cyto_dl_model.segmentation_plugin_model": [[5, "module-cyto_dl.api.cyto_dl_model.segmentation_plugin_model"]], "get_hardware_type() (cyto_dl.api.cyto_dl_model.segmentation_plugin_model.segmentationpluginmodel method)": [[5, "cyto_dl.api.cyto_dl_model.segmentation_plugin_model.SegmentationPluginModel.get_hardware_type"]], "get_input_channel() (cyto_dl.api.cyto_dl_model.segmentation_plugin_model.segmentationpluginmodel method)": [[5, "cyto_dl.api.cyto_dl_model.segmentation_plugin_model.SegmentationPluginModel.get_input_channel"]], "get_manifest_column_names() (cyto_dl.api.cyto_dl_model.segmentation_plugin_model.segmentationpluginmodel method)": [[5, "cyto_dl.api.cyto_dl_model.segmentation_plugin_model.SegmentationPluginModel.get_manifest_column_names"]], "get_patch_size() (cyto_dl.api.cyto_dl_model.segmentation_plugin_model.segmentationpluginmodel method)": [[5, "cyto_dl.api.cyto_dl_model.segmentation_plugin_model.SegmentationPluginModel.get_patch_size"]], "get_raw_image_channels() (cyto_dl.api.cyto_dl_model.segmentation_plugin_model.segmentationpluginmodel method)": [[5, "cyto_dl.api.cyto_dl_model.segmentation_plugin_model.SegmentationPluginModel.get_raw_image_channels"]], "get_split_column() (cyto_dl.api.cyto_dl_model.segmentation_plugin_model.segmentationpluginmodel method)": [[5, "cyto_dl.api.cyto_dl_model.segmentation_plugin_model.SegmentationPluginModel.get_split_column"]], "remove_split_column() (cyto_dl.api.cyto_dl_model.segmentation_plugin_model.segmentationpluginmodel method)": [[5, "cyto_dl.api.cyto_dl_model.segmentation_plugin_model.SegmentationPluginModel.remove_split_column"]], "set_hardware_type() (cyto_dl.api.cyto_dl_model.segmentation_plugin_model.segmentationpluginmodel method)": [[5, "cyto_dl.api.cyto_dl_model.segmentation_plugin_model.SegmentationPluginModel.set_hardware_type"]], "set_input_channel() (cyto_dl.api.cyto_dl_model.segmentation_plugin_model.segmentationpluginmodel method)": [[5, "cyto_dl.api.cyto_dl_model.segmentation_plugin_model.SegmentationPluginModel.set_input_channel"]], "set_manifest_column_names() (cyto_dl.api.cyto_dl_model.segmentation_plugin_model.segmentationpluginmodel method)": [[5, "cyto_dl.api.cyto_dl_model.segmentation_plugin_model.SegmentationPluginModel.set_manifest_column_names"]], "set_patch_size() (cyto_dl.api.cyto_dl_model.segmentation_plugin_model.segmentationpluginmodel method)": [[5, "cyto_dl.api.cyto_dl_model.segmentation_plugin_model.SegmentationPluginModel.set_patch_size"]], "set_raw_image_channels() (cyto_dl.api.cyto_dl_model.segmentation_plugin_model.segmentationpluginmodel method)": [[5, "cyto_dl.api.cyto_dl_model.segmentation_plugin_model.SegmentationPluginModel.set_raw_image_channels"]], "set_split_column() (cyto_dl.api.cyto_dl_model.segmentation_plugin_model.segmentationpluginmodel method)": [[5, "cyto_dl.api.cyto_dl_model.segmentation_plugin_model.SegmentationPluginModel.set_split_column"]], "autothresholdmethod (class in cyto_dl.api.data)": [[6, "cyto_dl.api.data.AutoThresholdMethod"]], "cpu (cyto_dl.api.data.hardwaretype attribute)": [[6, "cyto_dl.api.data.HardwareType.CPU"]], "experimenttype (class in cyto_dl.api.data)": [[6, "cyto_dl.api.data.ExperimentType"]], "gan (cyto_dl.api.data.experimenttype attribute)": [[6, "cyto_dl.api.data.ExperimentType.GAN"]], "gpu (cyto_dl.api.data.hardwaretype attribute)": [[6, "cyto_dl.api.data.HardwareType.GPU"]], "hardwaretype (class in cyto_dl.api.data)": [[6, "cyto_dl.api.data.HardwareType"]], "instance_seg (cyto_dl.api.data.experimenttype attribute)": [[6, "cyto_dl.api.data.ExperimentType.INSTANCE_SEG"]], "isodata (cyto_dl.api.data.autothresholdmethod attribute)": [[6, "cyto_dl.api.data.AutoThresholdMethod.ISODATA"]], "label_free (cyto_dl.api.data.experimenttype attribute)": [[6, "cyto_dl.api.data.ExperimentType.LABEL_FREE"]], "large (cyto_dl.api.data.patchsize attribute)": [[6, "cyto_dl.api.data.PatchSize.LARGE"]], "li (cyto_dl.api.data.autothresholdmethod attribute)": [[6, "cyto_dl.api.data.AutoThresholdMethod.LI"]], "local (cyto_dl.api.data.autothresholdmethod attribute)": [[6, "cyto_dl.api.data.AutoThresholdMethod.LOCAL"]], "mean (cyto_dl.api.data.autothresholdmethod attribute)": [[6, "cyto_dl.api.data.AutoThresholdMethod.MEAN"]], "medium (cyto_dl.api.data.patchsize attribute)": [[6, "cyto_dl.api.data.PatchSize.MEDIUM"]], "minimum (cyto_dl.api.data.autothresholdmethod attribute)": [[6, "cyto_dl.api.data.AutoThresholdMethod.MINIMUM"]], "multiotsu (cyto_dl.api.data.autothresholdmethod attribute)": [[6, "cyto_dl.api.data.AutoThresholdMethod.MULTIOTSU"]], "niblack (cyto_dl.api.data.autothresholdmethod attribute)": [[6, "cyto_dl.api.data.AutoThresholdMethod.NIBLACK"]], "otsu (cyto_dl.api.data.autothresholdmethod attribute)": [[6, "cyto_dl.api.data.AutoThresholdMethod.OTSU"]], "patchsize (class in cyto_dl.api.data)": [[6, "cyto_dl.api.data.PatchSize"]], "sauvola (cyto_dl.api.data.autothresholdmethod attribute)": [[6, "cyto_dl.api.data.AutoThresholdMethod.SAUVOLA"]], "segmentation (cyto_dl.api.data.experimenttype attribute)": [[6, "cyto_dl.api.data.ExperimentType.SEGMENTATION"]], "segmentation_array (cyto_dl.api.data.experimenttype attribute)": [[6, "cyto_dl.api.data.ExperimentType.SEGMENTATION_ARRAY"]], "segmentation_plugin (cyto_dl.api.data.experimenttype attribute)": [[6, "cyto_dl.api.data.ExperimentType.SEGMENTATION_PLUGIN"]], "small (cyto_dl.api.data.patchsize attribute)": [[6, "cyto_dl.api.data.PatchSize.SMALL"]], "triangle (cyto_dl.api.data.autothresholdmethod attribute)": [[6, "cyto_dl.api.data.AutoThresholdMethod.TRIANGLE"]], "yen (cyto_dl.api.data.autothresholdmethod attribute)": [[6, "cyto_dl.api.data.AutoThresholdMethod.YEN"]], "cyto_dl.api.data": [[6, "module-cyto_dl.api.data"]], "cytodlmodel (class in cyto_dl.api.model)": [[7, "cyto_dl.api.model.CytoDLModel"]], "cyto_dl.api.model": [[7, "module-cyto_dl.api.model"]], "download_example_data() (cyto_dl.api.model.cytodlmodel method)": [[7, "cyto_dl.api.model.CytoDLModel.download_example_data"]], "load_config_from_dict() (cyto_dl.api.model.cytodlmodel method)": [[7, "cyto_dl.api.model.CytoDLModel.load_config_from_dict"]], "load_config_from_file() (cyto_dl.api.model.cytodlmodel method)": [[7, "cyto_dl.api.model.CytoDLModel.load_config_from_file"]], "load_default_experiment() (cyto_dl.api.model.cytodlmodel method)": [[7, "cyto_dl.api.model.CytoDLModel.load_default_experiment"]], "override_config() (cyto_dl.api.model.cytodlmodel method)": [[7, "cyto_dl.api.model.CytoDLModel.override_config"]], "predict() (cyto_dl.api.model.cytodlmodel method)": [[7, "cyto_dl.api.model.CytoDLModel.predict"]], "print_config() (cyto_dl.api.model.cytodlmodel method)": [[7, "cyto_dl.api.model.CytoDLModel.print_config"]], "save_config() (cyto_dl.api.model.cytodlmodel method)": [[7, "cyto_dl.api.model.CytoDLModel.save_config"]], "train() (cyto_dl.api.model.cytodlmodel method)": [[7, "cyto_dl.api.model.CytoDLModel.train"]], "cyto_dl.callbacks": [[8, "module-cyto_dl.callbacks"]], "csvsaver (class in cyto_dl.callbacks.csv_saver)": [[10, "cyto_dl.callbacks.csv_saver.CSVSaver"]], "jepasaver (class in cyto_dl.callbacks.csv_saver)": [[10, "cyto_dl.callbacks.csv_saver.JEPASaver"]], "cyto_dl.callbacks.csv_saver": [[10, "module-cyto_dl.callbacks.csv_saver"]], "on_predict_epoch_end() (cyto_dl.callbacks.csv_saver.csvsaver method)": [[10, "cyto_dl.callbacks.csv_saver.CSVSaver.on_predict_epoch_end"]], "pred_to_csv() (cyto_dl.callbacks.csv_saver.csvsaver method)": [[10, "cyto_dl.callbacks.csv_saver.CSVSaver.pred_to_csv"]], "pred_to_csv() (cyto_dl.callbacks.csv_saver.jepasaver method)": [[10, "cyto_dl.callbacks.csv_saver.JEPASaver.pred_to_csv"]], "save_feats() (cyto_dl.callbacks.csv_saver.csvsaver method)": [[10, "cyto_dl.callbacks.csv_saver.CSVSaver.save_feats"]], "imagesaver (class in cyto_dl.callbacks.image_saver)": [[11, "cyto_dl.callbacks.image_saver.ImageSaver"]], "cyto_dl.callbacks.image_saver": [[11, "module-cyto_dl.callbacks.image_saver"]], "on_predict_batch_end() (cyto_dl.callbacks.image_saver.imagesaver method)": [[11, "cyto_dl.callbacks.image_saver.ImageSaver.on_predict_batch_end"]], "on_test_batch_end() (cyto_dl.callbacks.image_saver.imagesaver method)": [[11, "cyto_dl.callbacks.image_saver.ImageSaver.on_test_batch_end"]], "on_train_batch_end() (cyto_dl.callbacks.image_saver.imagesaver method)": [[11, "cyto_dl.callbacks.image_saver.ImageSaver.on_train_batch_end"]], "on_validation_batch_end() (cyto_dl.callbacks.image_saver.imagesaver method)": [[11, "cyto_dl.callbacks.image_saver.ImageSaver.on_validation_batch_end"]], "save() (cyto_dl.callbacks.image_saver.imagesaver method)": [[11, "cyto_dl.callbacks.image_saver.ImageSaver.save"]], "layerfreeze (class in cyto_dl.callbacks.layer_freeze)": [[13, "cyto_dl.callbacks.layer_freeze.LayerFreeze"]], "cyto_dl.callbacks.layer_freeze": [[13, "module-cyto_dl.callbacks.layer_freeze"]], "setup() (cyto_dl.callbacks.layer_freeze.layerfreeze method)": [[13, "cyto_dl.callbacks.layer_freeze.LayerFreeze.setup"]], "getembeddings (class in cyto_dl.callbacks.model_utils)": [[14, "cyto_dl.callbacks.model_utils.GetEmbeddings"]], "cyto_dl.callbacks.model_utils": [[14, "module-cyto_dl.callbacks.model_utils"]], "get_all_embeddings() (in module cyto_dl.callbacks.model_utils)": [[14, "cyto_dl.callbacks.model_utils.get_all_embeddings"]], "on_test_epoch_end() (cyto_dl.callbacks.model_utils.getembeddings method)": [[14, "cyto_dl.callbacks.model_utils.GetEmbeddings.on_test_epoch_end"]], "save_predictions_classifier() (in module cyto_dl.callbacks.model_utils)": [[14, "cyto_dl.callbacks.model_utils.save_predictions_classifier"]], "outlierdetection (class in cyto_dl.callbacks.outlier_detection)": [[15, "cyto_dl.callbacks.outlier_detection.OutlierDetection"]], "calculate_mahalanobis_hook() (cyto_dl.callbacks.outlier_detection.outlierdetection method)": [[15, "cyto_dl.callbacks.outlier_detection.OutlierDetection.calculate_mahalanobis_hook"]], "cyto_dl.callbacks.outlier_detection": [[15, "module-cyto_dl.callbacks.outlier_detection"]], "flatten_activations() (cyto_dl.callbacks.outlier_detection.outlierdetection method)": [[15, "cyto_dl.callbacks.outlier_detection.OutlierDetection.flatten_activations"]], "on_load_checkpoint() (cyto_dl.callbacks.outlier_detection.outlierdetection method)": [[15, "cyto_dl.callbacks.outlier_detection.OutlierDetection.on_load_checkpoint"]], "on_predict_batch_end() (cyto_dl.callbacks.outlier_detection.outlierdetection method)": [[15, "cyto_dl.callbacks.outlier_detection.OutlierDetection.on_predict_batch_end"]], "on_predict_epoch_start() (cyto_dl.callbacks.outlier_detection.outlierdetection method)": [[15, "cyto_dl.callbacks.outlier_detection.OutlierDetection.on_predict_epoch_start"]], "on_save_checkpoint() (cyto_dl.callbacks.outlier_detection.outlierdetection method)": [[15, "cyto_dl.callbacks.outlier_detection.OutlierDetection.on_save_checkpoint"]], "on_test_batch_end() (cyto_dl.callbacks.outlier_detection.outlierdetection method)": [[15, "cyto_dl.callbacks.outlier_detection.OutlierDetection.on_test_batch_end"]], "on_test_epoch_start() (cyto_dl.callbacks.outlier_detection.outlierdetection method)": [[15, "cyto_dl.callbacks.outlier_detection.OutlierDetection.on_test_epoch_start"]], "on_train_epoch_start() (cyto_dl.callbacks.outlier_detection.outlierdetection method)": [[15, "cyto_dl.callbacks.outlier_detection.OutlierDetection.on_train_epoch_start"]], "update_covariance_hook() (cyto_dl.callbacks.outlier_detection.outlierdetection method)": [[15, "cyto_dl.callbacks.outlier_detection.OutlierDetection.update_covariance_hook"]], "cyto_dl.dataframe": [[17, "module-cyto_dl.dataframe"]], "cyto_dl.dataframe.readers": [[18, "module-cyto_dl.dataframe.readers"]], "filter_columns() (in module cyto_dl.dataframe.readers)": [[18, "cyto_dl.dataframe.readers.filter_columns"]], "read_csv() (in module cyto_dl.dataframe.readers)": [[18, "cyto_dl.dataframe.readers.read_csv"]], "read_dataframe() (in module cyto_dl.dataframe.readers)": [[18, "cyto_dl.dataframe.readers.read_dataframe"]], "read_h5ad() (in module cyto_dl.dataframe.readers)": [[18, "cyto_dl.dataframe.readers.read_h5ad"]], "read_parquet() (in module cyto_dl.dataframe.readers)": [[18, "cyto_dl.dataframe.readers.read_parquet"]], "cyto_dl.dataframe.transforms": [[19, "module-cyto_dl.dataframe.transforms"]], "cyto_dl.dataframe.transforms.filter": [[20, "module-cyto_dl.dataframe.transforms.filter"]], "filter_columns() (in module cyto_dl.dataframe.transforms.filter)": [[20, "cyto_dl.dataframe.transforms.filter.filter_columns"]], "filter_rows() (in module cyto_dl.dataframe.transforms.filter)": [[20, "cyto_dl.dataframe.transforms.filter.filter_rows"]], "groupcols (class in cyto_dl.dataframe.transforms.group_cols)": [[21, "cyto_dl.dataframe.transforms.group_cols.GroupCols"]], "cyto_dl.dataframe.transforms.group_cols": [[21, "module-cyto_dl.dataframe.transforms.group_cols"]], "append_class_weights() (in module cyto_dl.dataframe.transforms.misc)": [[22, "cyto_dl.dataframe.transforms.misc.append_class_weights"]], "append_labels_to_integers() (in module cyto_dl.dataframe.transforms.misc)": [[22, "cyto_dl.dataframe.transforms.misc.append_labels_to_integers"]], "append_one_hot() (in module cyto_dl.dataframe.transforms.misc)": [[22, "cyto_dl.dataframe.transforms.misc.append_one_hot"]], "cyto_dl.dataframe.transforms.misc": [[22, "module-cyto_dl.dataframe.transforms.misc"]], "make_random_df() (in module cyto_dl.dataframe.transforms.misc)": [[22, "cyto_dl.dataframe.transforms.misc.make_random_df"]], "cyto_dl.dataframe.transforms.split": [[23, "module-cyto_dl.dataframe.transforms.split"]], "sample_n_each() (in module cyto_dl.dataframe.transforms.split)": [[23, "cyto_dl.dataframe.transforms.split.sample_n_each"]], "split_dataframe() (in module cyto_dl.dataframe.transforms.split)": [[23, "cyto_dl.dataframe.transforms.split.split_dataframe"]], "cyto_dl.datamodules": [[24, "module-cyto_dl.datamodules"]], "cyto_dl.datamodules.array": [[25, "module-cyto_dl.datamodules.array"]], "make_array_dataloader() (in module cyto_dl.datamodules.array)": [[25, "cyto_dl.datamodules.array.make_array_dataloader"]], "cyto_dl.datamodules.data_dict": [[26, "module-cyto_dl.datamodules.data_dict"]], "make_data_dict_dataloader() (in module cyto_dl.datamodules.data_dict)": [[26, "cyto_dl.datamodules.data_dict.make_data_dict_dataloader"]], "cyto_dl.datamodules.dataframe": [[27, "module-cyto_dl.datamodules.dataframe"]], "dataframedatamodule (class in cyto_dl.datamodules.dataframe.dataframe_datamodule)": [[28, "cyto_dl.datamodules.dataframe.dataframe_datamodule.DataframeDatamodule"]], "cyto_dl.datamodules.dataframe.dataframe_datamodule": [[28, "module-cyto_dl.datamodules.dataframe.dataframe_datamodule"]], "get_dataloader() (cyto_dl.datamodules.dataframe.dataframe_datamodule.dataframedatamodule method)": [[28, "cyto_dl.datamodules.dataframe.dataframe_datamodule.DataframeDatamodule.get_dataloader"]], "get_dataset() (cyto_dl.datamodules.dataframe.dataframe_datamodule.dataframedatamodule method)": [[28, "cyto_dl.datamodules.dataframe.dataframe_datamodule.DataframeDatamodule.get_dataset"]], "make_dataloader() (cyto_dl.datamodules.dataframe.dataframe_datamodule.dataframedatamodule method)": [[28, "cyto_dl.datamodules.dataframe.dataframe_datamodule.DataframeDatamodule.make_dataloader"]], "predict_dataloader() (cyto_dl.datamodules.dataframe.dataframe_datamodule.dataframedatamodule method)": [[28, "cyto_dl.datamodules.dataframe.dataframe_datamodule.DataframeDatamodule.predict_dataloader"]], "test_dataloader() (cyto_dl.datamodules.dataframe.dataframe_datamodule.dataframedatamodule method)": [[28, "cyto_dl.datamodules.dataframe.dataframe_datamodule.DataframeDatamodule.test_dataloader"]], "train_dataloader() (cyto_dl.datamodules.dataframe.dataframe_datamodule.dataframedatamodule method)": [[28, "cyto_dl.datamodules.dataframe.dataframe_datamodule.DataframeDatamodule.train_dataloader"]], "val_dataloader() (cyto_dl.datamodules.dataframe.dataframe_datamodule.dataframedatamodule method)": [[28, "cyto_dl.datamodules.dataframe.dataframe_datamodule.DataframeDatamodule.val_dataloader"]], "groupeddataframedatamodule (class in cyto_dl.datamodules.dataframe.grouped_dataframe_datamodule)": [[29, "cyto_dl.datamodules.dataframe.grouped_dataframe_datamodule.GroupedDataframeDatamodule"]], "cyto_dl.datamodules.dataframe.grouped_dataframe_datamodule": [[29, "module-cyto_dl.datamodules.dataframe.grouped_dataframe_datamodule"]], "make_dataloader() (cyto_dl.datamodules.dataframe.grouped_dataframe_datamodule.groupeddataframedatamodule method)": [[29, "cyto_dl.datamodules.dataframe.grouped_dataframe_datamodule.GroupedDataframeDatamodule.make_dataloader"]], "alternatingbatchsampler (class in cyto_dl.datamodules.dataframe.utils)": [[30, "cyto_dl.datamodules.dataframe.utils.AlternatingBatchSampler"]], "removenankeysd (class in cyto_dl.datamodules.dataframe.utils)": [[30, "cyto_dl.datamodules.dataframe.utils.RemoveNaNKeysd"]], "cyto_dl.datamodules.dataframe.utils": [[30, "module-cyto_dl.datamodules.dataframe.utils"]], "get_canonical_split_name() (in module cyto_dl.datamodules.dataframe.utils)": [[30, "cyto_dl.datamodules.dataframe.utils.get_canonical_split_name"]], "get_dataset() (in module cyto_dl.datamodules.dataframe.utils)": [[30, "cyto_dl.datamodules.dataframe.utils.get_dataset"]], "make_multiple_dataframe_splits() (in module cyto_dl.datamodules.dataframe.utils)": [[30, "cyto_dl.datamodules.dataframe.utils.make_multiple_dataframe_splits"]], "make_single_dataframe_splits() (in module cyto_dl.datamodules.dataframe.utils)": [[30, "cyto_dl.datamodules.dataframe.utils.make_single_dataframe_splits"]], "parse_transforms() (in module cyto_dl.datamodules.dataframe.utils)": [[30, "cyto_dl.datamodules.dataframe.utils.parse_transforms"]], "cyto_dl.datamodules.folder": [[31, "module-cyto_dl.datamodules.folder"]], "make_folder_dataloader() (in module cyto_dl.datamodules.folder)": [[31, "cyto_dl.datamodules.folder.make_folder_dataloader"]], "multidimimagedataset (class in cyto_dl.datamodules.multidim_image)": [[32, "cyto_dl.datamodules.multidim_image.MultiDimImageDataset"]], "create_metatensor() (cyto_dl.datamodules.multidim_image.multidimimagedataset method)": [[32, "cyto_dl.datamodules.multidim_image.MultiDimImageDataset.create_metatensor"]], "cyto_dl.datamodules.multidim_image": [[32, "module-cyto_dl.datamodules.multidim_image"]], "get_per_file_args() (cyto_dl.datamodules.multidim_image.multidimimagedataset method)": [[32, "cyto_dl.datamodules.multidim_image.MultiDimImageDataset.get_per_file_args"]], "is_batch() (cyto_dl.datamodules.multidim_image.multidimimagedataset method)": [[32, "cyto_dl.datamodules.multidim_image.MultiDimImageDataset.is_batch"]], "make_multidim_image_dataloader() (in module cyto_dl.datamodules.multidim_image)": [[32, "cyto_dl.datamodules.multidim_image.make_multidim_image_dataloader"]], "smartcachedatamodule (class in cyto_dl.datamodules.smartcache)": [[33, "cyto_dl.datamodules.smartcache.SmartcacheDatamodule"]], "cyto_dl.datamodules.smartcache": [[33, "module-cyto_dl.datamodules.smartcache"]], "get_per_file_args() (cyto_dl.datamodules.smartcache.smartcachedatamodule method)": [[33, "cyto_dl.datamodules.smartcache.SmartcacheDatamodule.get_per_file_args"]], "make_dataloader() (cyto_dl.datamodules.smartcache.smartcachedatamodule method)": [[33, "cyto_dl.datamodules.smartcache.SmartcacheDatamodule.make_dataloader"]], "predict_dataloader() (cyto_dl.datamodules.smartcache.smartcachedatamodule method)": [[33, "cyto_dl.datamodules.smartcache.SmartcacheDatamodule.predict_dataloader"]], "prepare_data() (cyto_dl.datamodules.smartcache.smartcachedatamodule method)": [[33, "cyto_dl.datamodules.smartcache.SmartcacheDatamodule.prepare_data"]], "setup() (cyto_dl.datamodules.smartcache.smartcachedatamodule method)": [[33, "cyto_dl.datamodules.smartcache.SmartcacheDatamodule.setup"]], "test_dataloader() (cyto_dl.datamodules.smartcache.smartcachedatamodule method)": [[33, "cyto_dl.datamodules.smartcache.SmartcacheDatamodule.test_dataloader"]], "train_dataloader() (cyto_dl.datamodules.smartcache.smartcachedatamodule method)": [[33, "cyto_dl.datamodules.smartcache.SmartcacheDatamodule.train_dataloader"]], "val_dataloader() (cyto_dl.datamodules.smartcache.smartcachedatamodule method)": [[33, "cyto_dl.datamodules.smartcache.SmartcacheDatamodule.val_dataloader"]], "cifar10 (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.CIFAR10"]], "cifar100 (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.CIFAR100"]], "clevrclassification (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.CLEVRClassification"]], "caltech101 (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.Caltech101"]], "caltech256 (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.Caltech256"]], "celeba (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.CelebA"]], "cityscapes (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.Cityscapes"]], "cococaptions (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.CocoCaptions"]], "cocodetection (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.CocoDetection"]], "country211 (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.Country211"]], "dtd (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.DTD"]], "datasetfolder (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.DatasetFolder"]], "emnist (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.EMNIST"]], "eurosat (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.EuroSAT"]], "fer2013 (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.FER2013"]], "fgvcaircraft (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.FGVCAircraft"]], "fakedata (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.FakeData"]], "fashionmnist (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.FashionMNIST"]], "flickr30k (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.Flickr30k"]], "flickr8k (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.Flickr8k"]], "flowers102 (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.Flowers102"]], "food101 (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.Food101"]], "gtsrb (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.GTSRB"]], "hmdb51 (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.HMDB51"]], "inaturalist (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.INaturalist"]], "imagefolder (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.ImageFolder"]], "imagenet (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.ImageNet"]], "kmnist (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.KMNIST"]], "kinetics (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.Kinetics"]], "kitti (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.Kitti"]], "lfwpairs (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.LFWPairs"]], "lfwpeople (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.LFWPeople"]], "lsun (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.LSUN"]], "lsunclass (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.LSUNClass"]], "mnist (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.MNIST"]], "movingmnist (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.MovingMNIST"]], "omniglot (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.Omniglot"]], "oxfordiiitpet (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.OxfordIIITPet"]], "pcam (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.PCAM"]], "phototour (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.PhotoTour"]], "places365 (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.Places365"]], "qmnist (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.QMNIST"]], "renderedsst2 (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.RenderedSST2"]], "sbdataset (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.SBDataset"]], "sbu (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.SBU"]], "semeion (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.SEMEION"]], "stl10 (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.STL10"]], "sun397 (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.SUN397"]], "svhn (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.SVHN"]], "stanfordcars (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.StanfordCars"]], "ucf101 (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.UCF101"]], "usps (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.USPS"]], "vocdetection (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.VOCDetection"]], "vocsegmentation (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.VOCSegmentation"]], "visiondataset (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.VisionDataset"]], "widerface (class in cyto_dl.datamodules.torchvision)": [[34, "cyto_dl.datamodules.torchvision.WIDERFace"]], "cyto_dl.datamodules.torchvision": [[34, "module-cyto_dl.datamodules.torchvision"]], "cyto_dl.eval": [[35, "module-cyto_dl.eval"]], "main() (in module cyto_dl.eval)": [[35, "cyto_dl.eval.main"]], "cyto_dl.image": [[36, "module-cyto_dl.image"]], "cyto_dl.image.io": [[37, "module-cyto_dl.image.io"]], "aicsimageloaderd (class in cyto_dl.image.io.aicsimage_loader)": [[38, "cyto_dl.image.io.aicsimage_loader.AICSImageLoaderd"]], "cyto_dl.image.io.aicsimage_loader": [[38, "module-cyto_dl.image.io.aicsimage_loader"]], "split_args() (cyto_dl.image.io.aicsimage_loader.aicsimageloaderd method)": [[38, "cyto_dl.image.io.aicsimage_loader.AICSImageLoaderd.split_args"]], "monaibioreader (class in cyto_dl.image.io.monai_bio_reader)": [[39, "cyto_dl.image.io.monai_bio_reader.MonaiBioReader"]], "cyto_dl.image.io.monai_bio_reader": [[39, "module-cyto_dl.image.io.monai_bio_reader"]], "get_data() (cyto_dl.image.io.monai_bio_reader.monaibioreader method)": [[39, "cyto_dl.image.io.monai_bio_reader.MonaiBioReader.get_data"]], "read() (cyto_dl.image.io.monai_bio_reader.monaibioreader method)": [[39, "cyto_dl.image.io.monai_bio_reader.MonaiBioReader.read"]], "verify_suffix() (cyto_dl.image.io.monai_bio_reader.monaibioreader method)": [[39, "cyto_dl.image.io.monai_bio_reader.MonaiBioReader.verify_suffix"]], "readnumpyfile (class in cyto_dl.image.io.numpy_reader)": [[40, "cyto_dl.image.io.numpy_reader.ReadNumpyFile"]], "cyto_dl.image.io.numpy_reader": [[40, "module-cyto_dl.image.io.numpy_reader"]], "omezarrreader (class in cyto_dl.image.io.ome_zarr_reader)": [[41, "cyto_dl.image.io.ome_zarr_reader.OmeZarrReader"]], "cyto_dl.image.io.ome_zarr_reader": [[41, "module-cyto_dl.image.io.ome_zarr_reader"]], "get_data() (cyto_dl.image.io.ome_zarr_reader.omezarrreader method)": [[41, "cyto_dl.image.io.ome_zarr_reader.OmeZarrReader.get_data"]], "read() (cyto_dl.image.io.ome_zarr_reader.omezarrreader method)": [[41, "cyto_dl.image.io.ome_zarr_reader.OmeZarrReader.read"]], "verify_suffix() (cyto_dl.image.io.ome_zarr_reader.omezarrreader method)": [[41, "cyto_dl.image.io.ome_zarr_reader.OmeZarrReader.verify_suffix"]], "polygonloaderd (class in cyto_dl.image.io.polygon_loader)": [[42, "cyto_dl.image.io.polygon_loader.PolygonLoaderd"]], "cyto_dl.image.io.polygon_loader": [[42, "module-cyto_dl.image.io.polygon_loader"]], "skimagereader (class in cyto_dl.image.io.skimage_reader)": [[43, "cyto_dl.image.io.skimage_reader.SkimageReader"]], "cyto_dl.image.io.skimage_reader": [[43, "module-cyto_dl.image.io.skimage_reader"]], "get_data() (cyto_dl.image.io.skimage_reader.skimagereader method)": [[43, "cyto_dl.image.io.skimage_reader.SkimageReader.get_data"]], "read() (cyto_dl.image.io.skimage_reader.skimagereader method)": [[43, "cyto_dl.image.io.skimage_reader.SkimageReader.read"]], "verify_suffix() (cyto_dl.image.io.skimage_reader.skimagereader method)": [[43, "cyto_dl.image.io.skimage_reader.SkimageReader.verify_suffix"]], "cyto_dl.image.transforms": [[44, "module-cyto_dl.image.transforms"]], "addmeta (class in cyto_dl.image.transforms.add_meta)": [[45, "cyto_dl.image.transforms.add_meta.AddMeta"]], "cyto_dl.image.transforms.add_meta": [[45, "module-cyto_dl.image.transforms.add_meta"]], "brightsampler (class in cyto_dl.image.transforms.bright_sampler)": [[46, "cyto_dl.image.transforms.bright_sampler.BrightSampler"]], "cyto_dl.image.transforms.bright_sampler": [[46, "module-cyto_dl.image.transforms.bright_sampler"]], "clip (class in cyto_dl.image.transforms.clip)": [[47, "cyto_dl.image.transforms.clip.Clip"]], "clipd (class in cyto_dl.image.transforms.clip)": [[47, "cyto_dl.image.transforms.clip.Clipd"]], "cyto_dl.image.transforms.clip": [[47, "module-cyto_dl.image.transforms.clip"]], "contrastadjust (class in cyto_dl.image.transforms.contrastadjust)": [[48, "cyto_dl.image.transforms.contrastadjust.ContrastAdjust"]], "contrastadjustd (class in cyto_dl.image.transforms.contrastadjust)": [[48, "cyto_dl.image.transforms.contrastadjust.ContrastAdjustd"]], "cyto_dl.image.transforms.contrastadjust": [[48, "module-cyto_dl.image.transforms.contrastadjust"]], "jepamaskgenerator (class in cyto_dl.image.transforms.generate_jepa_masks)": [[49, "cyto_dl.image.transforms.generate_jepa_masks.JEPAMaskGenerator"]], "cyto_dl.image.transforms.generate_jepa_masks": [[49, "module-cyto_dl.image.transforms.generate_jepa_masks"]], "remove_excess_pixels() (cyto_dl.image.transforms.generate_jepa_masks.jepamaskgenerator method)": [[49, "cyto_dl.image.transforms.generate_jepa_masks.JEPAMaskGenerator.remove_excess_pixels"]], "merged (class in cyto_dl.image.transforms.merge)": [[50, "cyto_dl.image.transforms.merge.Merged"]], "cyto_dl.image.transforms.merge": [[50, "module-cyto_dl.image.transforms.merge"]], "randommultiscalecropd (class in cyto_dl.image.transforms.multiscale_cropper)": [[51, "cyto_dl.image.transforms.multiscale_cropper.RandomMultiScaleCropd"]], "cyto_dl.image.transforms.multiscale_cropper": [[51, "module-cyto_dl.image.transforms.multiscale_cropper"]], "generate_slices() (cyto_dl.image.transforms.multiscale_cropper.randommultiscalecropd method)": [[51, "cyto_dl.image.transforms.multiscale_cropper.RandomMultiScaleCropd.generate_slices"]], "padzd (class in cyto_dl.image.transforms.pad)": [[52, "cyto_dl.image.transforms.pad.PadZd"]], "cyto_dl.image.transforms.pad": [[52, "module-cyto_dl.image.transforms.pad"]], "maxprojectd (class in cyto_dl.image.transforms.project)": [[53, "cyto_dl.image.transforms.project.MaxProjectd"]], "cyto_dl.image.transforms.project": [[53, "module-cyto_dl.image.transforms.project"]], "save (class in cyto_dl.image.transforms.save)": [[55, "cyto_dl.image.transforms.save.Save"]], "saved (class in cyto_dl.image.transforms.save)": [[55, "cyto_dl.image.transforms.save.Saved"]], "cyto_dl.image.transforms.save": [[55, "module-cyto_dl.image.transforms.save"]], "cropresize (class in cyto_dl.image.transforms.track_transforms)": [[57, "cyto_dl.image.transforms.track_transforms.CropResize"]], "generatetracklabels (class in cyto_dl.image.transforms.track_transforms)": [[57, "cyto_dl.image.transforms.track_transforms.GenerateTrackLabels"]], "perchannel (class in cyto_dl.image.transforms.track_transforms)": [[57, "cyto_dl.image.transforms.track_transforms.PerChannel"]], "splittrackd (class in cyto_dl.image.transforms.track_transforms)": [[57, "cyto_dl.image.transforms.track_transforms.SplitTrackd"]], "cyto_dl.image.transforms.track_transforms": [[57, "module-cyto_dl.image.transforms.track_transforms"]], "cyto_dl.loggers": [[58, "module-cyto_dl.loggers"]], "mlflowlogger (class in cyto_dl.loggers.mlflow)": [[59, "cyto_dl.loggers.mlflow.MLFlowLogger"]], "after_save_checkpoint() (cyto_dl.loggers.mlflow.mlflowlogger method)": [[59, "cyto_dl.loggers.mlflow.MLFlowLogger.after_save_checkpoint"]], "cyto_dl.loggers.mlflow": [[59, "module-cyto_dl.loggers.mlflow"]], "log_hyperparams() (cyto_dl.loggers.mlflow.mlflowlogger method)": [[59, "cyto_dl.loggers.mlflow.MLFlowLogger.log_hyperparams"]], "log_metrics() (cyto_dl.loggers.mlflow.mlflowlogger method)": [[59, "cyto_dl.loggers.mlflow.MLFlowLogger.log_metrics"]], "cyto_dl.models": [[60, "module-cyto_dl.models"]], "basemodel (class in cyto_dl.models.base_model)": [[61, "cyto_dl.models.base_model.BaseModel"]], "basemodelmeta (class in cyto_dl.models.base_model)": [[61, "cyto_dl.models.base_model.BaseModelMeta"]], "compute_metrics() (cyto_dl.models.base_model.basemodel method)": [[61, "cyto_dl.models.base_model.BaseModel.compute_metrics"]], "configure_optimizers() (cyto_dl.models.base_model.basemodel method)": [[61, "cyto_dl.models.base_model.BaseModel.configure_optimizers"]], "cyto_dl.models.base_model": [[61, "module-cyto_dl.models.base_model"]], "forward() (cyto_dl.models.base_model.basemodel method)": [[61, "cyto_dl.models.base_model.BaseModel.forward"]], "model_step() (cyto_dl.models.base_model.basemodel method)": [[61, "cyto_dl.models.base_model.BaseModel.model_step"]], "on_train_start() (cyto_dl.models.base_model.basemodel method)": [[61, "cyto_dl.models.base_model.BaseModel.on_train_start"]], "parse_batch() (cyto_dl.models.base_model.basemodel method)": [[61, "cyto_dl.models.base_model.BaseModel.parse_batch"]], "predict_step() (cyto_dl.models.base_model.basemodel method)": [[61, "cyto_dl.models.base_model.BaseModel.predict_step"]], "test_step() (cyto_dl.models.base_model.basemodel method)": [[61, "cyto_dl.models.base_model.BaseModel.test_step"]], "training_step() (cyto_dl.models.base_model.basemodel method)": [[61, "cyto_dl.models.base_model.BaseModel.training_step"]], "validation_step() (cyto_dl.models.base_model.basemodel method)": [[61, "cyto_dl.models.base_model.BaseModel.validation_step"]], "basicmodel (class in cyto_dl.models.basic_model)": [[62, "cyto_dl.models.basic_model.BasicModel"]], "cyto_dl.models.basic_model": [[62, "module-cyto_dl.models.basic_model"]], "forward() (cyto_dl.models.basic_model.basicmodel method)": [[62, "cyto_dl.models.basic_model.BasicModel.forward"]], "model_step() (cyto_dl.models.basic_model.basicmodel method)": [[62, "cyto_dl.models.basic_model.BasicModel.model_step"]], "cyto_dl.models.classification": [[63, "module-cyto_dl.models.classification"]], "classification (class in cyto_dl.models.classification.classification)": [[64, "cyto_dl.models.classification.classification.Classification"]], "cyto_dl.models.classification.classification": [[64, "module-cyto_dl.models.classification.classification"]], "forward() (cyto_dl.models.classification.classification.classification method)": [[64, "cyto_dl.models.classification.classification.Classification.forward"]], "model_step() (cyto_dl.models.classification.classification.classification method)": [[64, "cyto_dl.models.classification.classification.Classification.model_step"]], "predict_step() (cyto_dl.models.classification.classification.classification method)": [[64, "cyto_dl.models.classification.classification.Classification.predict_step"]], "save_images() (cyto_dl.models.classification.classification.classification method)": [[64, "cyto_dl.models.classification.classification.Classification.save_images"]], "should_save_image() (cyto_dl.models.classification.classification.classification method)": [[64, "cyto_dl.models.classification.classification.Classification.should_save_image"]], "timepointclassification (class in cyto_dl.models.classification.timepoint_classification)": [[65, "cyto_dl.models.classification.timepoint_classification.TimepointClassification"]], "cyto_dl.models.classification.timepoint_classification": [[65, "module-cyto_dl.models.classification.timepoint_classification"]], "predict_step() (cyto_dl.models.classification.timepoint_classification.timepointclassification method)": [[65, "cyto_dl.models.classification.timepoint_classification.TimepointClassification.predict_step"]], "cyto_dl.models.contrastive": [[66, "module-cyto_dl.models.contrastive"]], "contrastive (class in cyto_dl.models.contrastive.contrastive)": [[67, "cyto_dl.models.contrastive.contrastive.Contrastive"]], "cyto_dl.models.contrastive.contrastive": [[67, "module-cyto_dl.models.contrastive.contrastive"]], "forward() (cyto_dl.models.contrastive.contrastive.contrastive method)": [[67, "cyto_dl.models.contrastive.contrastive.Contrastive.forward"]], "model_step() (cyto_dl.models.contrastive.contrastive.contrastive method)": [[67, "cyto_dl.models.contrastive.contrastive.Contrastive.model_step"]], "plot_classes() (cyto_dl.models.contrastive.contrastive.contrastive method)": [[67, "cyto_dl.models.contrastive.contrastive.Contrastive.plot_classes"]], "plot_neighbors() (cyto_dl.models.contrastive.contrastive.contrastive method)": [[67, "cyto_dl.models.contrastive.contrastive.Contrastive.plot_neighbors"]], "predict_step() (cyto_dl.models.contrastive.contrastive.contrastive method)": [[67, "cyto_dl.models.contrastive.contrastive.Contrastive.predict_step"]], "cyto_dl.models.im2im": [[72, "module-cyto_dl.models.im2im"]], "gan (class in cyto_dl.models.im2im.gan)": [[73, "cyto_dl.models.im2im.gan.GAN"]], "configure_optimizers() (cyto_dl.models.im2im.gan.gan method)": [[73, "cyto_dl.models.im2im.gan.GAN.configure_optimizers"]], "cyto_dl.models.im2im.gan": [[73, "module-cyto_dl.models.im2im.gan"]], "model_step() (cyto_dl.models.im2im.gan.gan method)": [[73, "cyto_dl.models.im2im.gan.GAN.model_step"]], "predict_step() (cyto_dl.models.im2im.gan.gan method)": [[73, "cyto_dl.models.im2im.gan.GAN.predict_step"]], "multitaskim2im (class in cyto_dl.models.im2im.multi_task)": [[74, "cyto_dl.models.im2im.multi_task.MultiTaskIm2Im"]], "configure_optimizers() (cyto_dl.models.im2im.multi_task.multitaskim2im method)": [[74, "cyto_dl.models.im2im.multi_task.MultiTaskIm2Im.configure_optimizers"]], "cyto_dl.models.im2im.multi_task": [[74, "module-cyto_dl.models.im2im.multi_task"]], "forward() (cyto_dl.models.im2im.multi_task.multitaskim2im method)": [[74, "cyto_dl.models.im2im.multi_task.MultiTaskIm2Im.forward"]], "get_n_postprocess_image() (cyto_dl.models.im2im.multi_task.multitaskim2im method)": [[74, "cyto_dl.models.im2im.multi_task.MultiTaskIm2Im.get_n_postprocess_image"]], "get_per_head() (cyto_dl.models.im2im.multi_task.multitaskim2im method)": [[74, "cyto_dl.models.im2im.multi_task.MultiTaskIm2Im.get_per_head"]], "model_step() (cyto_dl.models.im2im.multi_task.multitaskim2im method)": [[74, "cyto_dl.models.im2im.multi_task.MultiTaskIm2Im.model_step"]], "predict_step() (cyto_dl.models.im2im.multi_task.multitaskim2im method)": [[74, "cyto_dl.models.im2im.multi_task.MultiTaskIm2Im.predict_step"]], "run_forward() (cyto_dl.models.im2im.multi_task.multitaskim2im method)": [[74, "cyto_dl.models.im2im.multi_task.MultiTaskIm2Im.run_forward"]], "test_step() (cyto_dl.models.im2im.multi_task.multitaskim2im method)": [[74, "cyto_dl.models.im2im.multi_task.MultiTaskIm2Im.test_step"]], "training_step() (cyto_dl.models.im2im.multi_task.multitaskim2im method)": [[74, "cyto_dl.models.im2im.multi_task.MultiTaskIm2Im.training_step"]], "validation_step() (cyto_dl.models.im2im.multi_task.multitaskim2im method)": [[74, "cyto_dl.models.im2im.multi_task.MultiTaskIm2Im.validation_step"]], "cyto_dl.models.im2im.utils": [[75, "module-cyto_dl.models.im2im.utils"]], "instancesegcluster (class in cyto_dl.models.im2im.utils.instance_seg)": [[76, "cyto_dl.models.im2im.utils.instance_seg.InstanceSegCluster"]], "instancesegloss (class in cyto_dl.models.im2im.utils.instance_seg)": [[76, "cyto_dl.models.im2im.utils.instance_seg.InstanceSegLoss"]], "instancesegpreprocessd (class in cyto_dl.models.im2im.utils.instance_seg)": [[76, "cyto_dl.models.im2im.utils.instance_seg.InstanceSegPreprocessd"]], "instancesegrandflipd (class in cyto_dl.models.im2im.utils.instance_seg)": [[76, "cyto_dl.models.im2im.utils.instance_seg.InstanceSegRandFlipd"]], "cluster_object() (cyto_dl.models.im2im.utils.instance_seg.instancesegcluster method)": [[76, "cyto_dl.models.im2im.utils.instance_seg.InstanceSegCluster.cluster_object"]], "cyto_dl.models.im2im.utils.instance_seg": [[76, "module-cyto_dl.models.im2im.utils.instance_seg"]], "embed_from_skel() (cyto_dl.models.im2im.utils.instance_seg.instancesegpreprocessd method)": [[76, "cyto_dl.models.im2im.utils.instance_seg.InstanceSegPreprocessd.embed_from_skel"]], "kd_clustering() (cyto_dl.models.im2im.utils.instance_seg.instancesegcluster method)": [[76, "cyto_dl.models.im2im.utils.instance_seg.InstanceSegCluster.kd_clustering"]], "keep_largest_cc() (cyto_dl.models.im2im.utils.instance_seg.instancesegpreprocessd method)": [[76, "cyto_dl.models.im2im.utils.instance_seg.InstanceSegPreprocessd.keep_largest_cc"]], "label_2d() (cyto_dl.models.im2im.utils.instance_seg.instancesegpreprocessd method)": [[76, "cyto_dl.models.im2im.utils.instance_seg.InstanceSegPreprocessd.label_2d"]], "pad_slice() (in module cyto_dl.models.im2im.utils.instance_seg)": [[76, "cyto_dl.models.im2im.utils.instance_seg.pad_slice"]], "remove_small_skeletons() (cyto_dl.models.im2im.utils.instance_seg.instancesegcluster method)": [[76, "cyto_dl.models.im2im.utils.instance_seg.InstanceSegCluster.remove_small_skeletons"]], "shrink() (cyto_dl.models.im2im.utils.instance_seg.instancesegpreprocessd method)": [[76, "cyto_dl.models.im2im.utils.instance_seg.InstanceSegPreprocessd.shrink"]], "skeleton_tall() (cyto_dl.models.im2im.utils.instance_seg.instancesegpreprocessd method)": [[76, "cyto_dl.models.im2im.utils.instance_seg.InstanceSegPreprocessd.skeleton_tall"]], "smooth_embedding() (cyto_dl.models.im2im.utils.instance_seg.instancesegpreprocessd method)": [[76, "cyto_dl.models.im2im.utils.instance_seg.InstanceSegPreprocessd.smooth_embedding"]], "topology_preserving_thinning() (cyto_dl.models.im2im.utils.instance_seg.instancesegpreprocessd method)": [[76, "cyto_dl.models.im2im.utils.instance_seg.InstanceSegPreprocessd.topology_preserving_thinning"]], "noiseannealer (class in cyto_dl.models.im2im.utils.noise_annealer)": [[77, "cyto_dl.models.im2im.utils.noise_annealer.NoiseAnnealer"]], "cyto_dl.models.im2im.utils.noise_annealer": [[77, "module-cyto_dl.models.im2im.utils.noise_annealer"]], "update_noise() (cyto_dl.models.im2im.utils.noise_annealer.noiseannealer method)": [[77, "cyto_dl.models.im2im.utils.noise_annealer.NoiseAnnealer.update_noise"]], "cyto_dl.models.im2im.utils.postprocessing": [[78, "module-cyto_dl.models.im2im.utils.postprocessing"]], "actthreshlabel (class in cyto_dl.models.im2im.utils.postprocessing.act_thresh_label)": [[79, "cyto_dl.models.im2im.utils.postprocessing.act_thresh_label.ActThreshLabel"]], "cyto_dl.models.im2im.utils.postprocessing.act_thresh_label": [[79, "module-cyto_dl.models.im2im.utils.postprocessing.act_thresh_label"]], "autothreshold (class in cyto_dl.models.im2im.utils.postprocessing.auto_thresh)": [[80, "cyto_dl.models.im2im.utils.postprocessing.auto_thresh.AutoThreshold"]], "cyto_dl.models.im2im.utils.postprocessing.auto_thresh": [[80, "module-cyto_dl.models.im2im.utils.postprocessing.auto_thresh"]], "dicttoim (class in cyto_dl.models.im2im.utils.postprocessing.dict_to_im)": [[81, "cyto_dl.models.im2im.utils.postprocessing.dict_to_im.DictToIm"]], "cyto_dl.models.im2im.utils.postprocessing.dict_to_im": [[81, "module-cyto_dl.models.im2im.utils.postprocessing.dict_to_im"]], "detach() (in module cyto_dl.models.im2im.utils.postprocessing.dict_to_im)": [[81, "cyto_dl.models.im2im.utils.postprocessing.dict_to_im.detach"]], "cyto_dl.models.jepa": [[82, "module-cyto_dl.models.jepa"]], "ijepa (class in cyto_dl.models.jepa.ijepa)": [[83, "cyto_dl.models.jepa.ijepa.IJEPA"]], "cyto_dl.models.jepa.ijepa": [[83, "module-cyto_dl.models.jepa.ijepa"]], "model_step() (cyto_dl.models.jepa.ijepa.ijepa method)": [[83, "cyto_dl.models.jepa.ijepa.IJEPA.model_step"]], "iwm (class in cyto_dl.models.jepa.iwm)": [[84, "cyto_dl.models.jepa.iwm.IWM"]], "cyto_dl.models.jepa.iwm": [[84, "module-cyto_dl.models.jepa.iwm"]], "extract_embeddings() (cyto_dl.models.jepa.iwm.iwm method)": [[84, "cyto_dl.models.jepa.iwm.IWM.extract_embeddings"]], "get_predict_masks() (cyto_dl.models.jepa.iwm.iwm method)": [[84, "cyto_dl.models.jepa.iwm.IWM.get_predict_masks"]], "model_step() (cyto_dl.models.jepa.iwm.iwm method)": [[84, "cyto_dl.models.jepa.iwm.IWM.model_step"]], "predict_step() (cyto_dl.models.jepa.iwm.iwm method)": [[84, "cyto_dl.models.jepa.iwm.IWM.predict_step"]], "test_step() (cyto_dl.models.jepa.iwm.iwm method)": [[84, "cyto_dl.models.jepa.iwm.IWM.test_step"]], "jepabase (class in cyto_dl.models.jepa.jepa_base)": [[85, "cyto_dl.models.jepa.jepa_base.JEPABase"]], "configure_optimizers() (cyto_dl.models.jepa.jepa_base.jepabase method)": [[85, "cyto_dl.models.jepa.jepa_base.JEPABase.configure_optimizers"]], "cyto_dl.models.jepa.jepa_base": [[85, "module-cyto_dl.models.jepa.jepa_base"]], "forward() (cyto_dl.models.jepa.jepa_base.jepabase method)": [[85, "cyto_dl.models.jepa.jepa_base.JEPABase.forward"]], "get_context_embeddings() (cyto_dl.models.jepa.jepa_base.jepabase method)": [[85, "cyto_dl.models.jepa.jepa_base.JEPABase.get_context_embeddings"]], "get_mask() (cyto_dl.models.jepa.jepa_base.jepabase method)": [[85, "cyto_dl.models.jepa.jepa_base.JEPABase.get_mask"]], "get_target_embeddings() (cyto_dl.models.jepa.jepa_base.jepabase method)": [[85, "cyto_dl.models.jepa.jepa_base.JEPABase.get_target_embeddings"]], "model_step() (cyto_dl.models.jepa.jepa_base.jepabase method)": [[85, "cyto_dl.models.jepa.jepa_base.JEPABase.model_step"]], "predict_step() (cyto_dl.models.jepa.jepa_base.jepabase method)": [[85, "cyto_dl.models.jepa.jepa_base.JEPABase.predict_step"]], "remove_first_dim() (cyto_dl.models.jepa.jepa_base.jepabase method)": [[85, "cyto_dl.models.jepa.jepa_base.JEPABase.remove_first_dim"]], "update_teacher() (cyto_dl.models.jepa.jepa_base.jepabase method)": [[85, "cyto_dl.models.jepa.jepa_base.JEPABase.update_teacher"]], "cyto_dl.models.utils": [[86, "module-cyto_dl.models.utils"]], "cyto_dl.models.utils.mlflow": [[87, "module-cyto_dl.models.utils.mlflow"]], "get_config() (in module cyto_dl.models.utils.mlflow)": [[87, "cyto_dl.models.utils.mlflow.get_config"]], "load_model_from_checkpoint() (in module cyto_dl.models.utils.mlflow)": [[87, "cyto_dl.models.utils.mlflow.load_model_from_checkpoint"]], "cyto_dl.models.utils.utils": [[88, "module-cyto_dl.models.utils.utils"]], "find_indices() (in module cyto_dl.models.utils.utils)": [[88, "cyto_dl.models.utils.utils.find_indices"]], "basevae (class in cyto_dl.models.vae.base_vae)": [[90, "cyto_dl.models.vae.base_vae.BaseVAE"]], "calculate_elbo() (cyto_dl.models.vae.base_vae.basevae method)": [[90, "cyto_dl.models.vae.base_vae.BaseVAE.calculate_elbo"]], "calculate_rcl() (cyto_dl.models.vae.base_vae.basevae method)": [[90, "cyto_dl.models.vae.base_vae.BaseVAE.calculate_rcl"]], "calculate_rcl_dict() (cyto_dl.models.vae.base_vae.basevae method)": [[90, "cyto_dl.models.vae.base_vae.BaseVAE.calculate_rcl_dict"]], "cyto_dl.models.vae.base_vae": [[90, "module-cyto_dl.models.vae.base_vae"]], "decode() (cyto_dl.models.vae.base_vae.basevae method)": [[90, "cyto_dl.models.vae.base_vae.BaseVAE.decode"]], "encode() (cyto_dl.models.vae.base_vae.basevae method)": [[90, "cyto_dl.models.vae.base_vae.BaseVAE.encode"]], "forward() (cyto_dl.models.vae.base_vae.basevae method)": [[90, "cyto_dl.models.vae.base_vae.BaseVAE.forward"]], "model_step() (cyto_dl.models.vae.base_vae.basevae method)": [[90, "cyto_dl.models.vae.base_vae.BaseVAE.model_step"]], "sample_z() (cyto_dl.models.vae.base_vae.basevae method)": [[90, "cyto_dl.models.vae.base_vae.BaseVAE.sample_z"]], "cyto_dl.models.vae.priors": [[99, "module-cyto_dl.models.vae.priors"]], "prior (class in cyto_dl.models.vae.priors.abstract_prior)": [[100, "cyto_dl.models.vae.priors.abstract_prior.Prior"]], "cyto_dl.models.vae.priors.abstract_prior": [[100, "module-cyto_dl.models.vae.priors.abstract_prior"]], "forward() (cyto_dl.models.vae.priors.abstract_prior.prior method)": [[100, "cyto_dl.models.vae.priors.abstract_prior.Prior.forward"]], "kl_divergence() (cyto_dl.models.vae.priors.abstract_prior.prior method)": [[100, "cyto_dl.models.vae.priors.abstract_prior.Prior.kl_divergence"]], "param_size (cyto_dl.models.vae.priors.abstract_prior.prior property)": [[100, "cyto_dl.models.vae.priors.abstract_prior.Prior.param_size"]], "sample() (cyto_dl.models.vae.priors.abstract_prior.prior method)": [[100, "cyto_dl.models.vae.priors.abstract_prior.Prior.sample"]], "diagonalgaussianprior (class in cyto_dl.models.vae.priors.gaussian)": [[101, "cyto_dl.models.vae.priors.gaussian.DiagonalGaussianPrior"]], "isotropicgaussianprior (class in cyto_dl.models.vae.priors.gaussian)": [[101, "cyto_dl.models.vae.priors.gaussian.IsotropicGaussianPrior"]], "compute_tc_penalty() (in module cyto_dl.models.vae.priors.gaussian)": [[101, "cyto_dl.models.vae.priors.gaussian.compute_tc_penalty"]], "cyto_dl.models.vae.priors.gaussian": [[101, "module-cyto_dl.models.vae.priors.gaussian"]], "forward() (cyto_dl.models.vae.priors.gaussian.diagonalgaussianprior method)": [[101, "cyto_dl.models.vae.priors.gaussian.DiagonalGaussianPrior.forward"]], "forward() (cyto_dl.models.vae.priors.gaussian.isotropicgaussianprior method)": [[101, "cyto_dl.models.vae.priors.gaussian.IsotropicGaussianPrior.forward"]], "kl_divergence() (cyto_dl.models.vae.priors.gaussian.diagonalgaussianprior class method)": [[101, "cyto_dl.models.vae.priors.gaussian.DiagonalGaussianPrior.kl_divergence"]], "kl_divergence() (cyto_dl.models.vae.priors.gaussian.isotropicgaussianprior class method)": [[101, "cyto_dl.models.vae.priors.gaussian.IsotropicGaussianPrior.kl_divergence"]], "param_size (cyto_dl.models.vae.priors.gaussian.diagonalgaussianprior property)": [[101, "cyto_dl.models.vae.priors.gaussian.DiagonalGaussianPrior.param_size"]], "param_size (cyto_dl.models.vae.priors.gaussian.isotropicgaussianprior property)": [[101, "cyto_dl.models.vae.priors.gaussian.IsotropicGaussianPrior.param_size"]], "sample() (cyto_dl.models.vae.priors.gaussian.isotropicgaussianprior class method)": [[101, "cyto_dl.models.vae.priors.gaussian.IsotropicGaussianPrior.sample"]], "identityprior (class in cyto_dl.models.vae.priors.identity_prior)": [[102, "cyto_dl.models.vae.priors.identity_prior.IdentityPrior"]], "cyto_dl.models.vae.priors.identity_prior": [[102, "module-cyto_dl.models.vae.priors.identity_prior"]], "forward() (cyto_dl.models.vae.priors.identity_prior.identityprior method)": [[102, "cyto_dl.models.vae.priors.identity_prior.IdentityPrior.forward"]], "param_size (cyto_dl.models.vae.priors.identity_prior.identityprior property)": [[102, "cyto_dl.models.vae.priors.identity_prior.IdentityPrior.param_size"]], "jointprior (class in cyto_dl.models.vae.priors.joint_prior)": [[103, "cyto_dl.models.vae.priors.joint_prior.JointPrior"]], "cyto_dl.models.vae.priors.joint_prior": [[103, "module-cyto_dl.models.vae.priors.joint_prior"]], "forward() (cyto_dl.models.vae.priors.joint_prior.jointprior method)": [[103, "cyto_dl.models.vae.priors.joint_prior.JointPrior.forward"]], "kl_divergence() (cyto_dl.models.vae.priors.joint_prior.jointprior method)": [[103, "cyto_dl.models.vae.priors.joint_prior.JointPrior.kl_divergence"]], "param_size (cyto_dl.models.vae.priors.joint_prior.jointprior property)": [[103, "cyto_dl.models.vae.priors.joint_prior.JointPrior.param_size"]], "sample() (cyto_dl.models.vae.priors.joint_prior.jointprior method)": [[103, "cyto_dl.models.vae.priors.joint_prior.JointPrior.sample"]], "cyto_dl.nn": [[105, "module-cyto_dl.nn"]], "cyto_dl.nn.discriminators": [[106, "module-cyto_dl.nn.discriminators"]], "multiscalediscriminator (class in cyto_dl.nn.discriminators.multi_scale_discriminator)": [[107, "cyto_dl.nn.discriminators.multi_scale_discriminator.MultiScaleDiscriminator"]], "cyto_dl.nn.discriminators.multi_scale_discriminator": [[107, "module-cyto_dl.nn.discriminators.multi_scale_discriminator"]], "forward() (cyto_dl.nn.discriminators.multi_scale_discriminator.multiscalediscriminator method)": [[107, "cyto_dl.nn.discriminators.multi_scale_discriminator.MultiScaleDiscriminator.forward"]], "nlayerdiscriminator (class in cyto_dl.nn.discriminators.n_layer_discriminator)": [[108, "cyto_dl.nn.discriminators.n_layer_discriminator.NLayerDiscriminator"]], "cyto_dl.nn.discriminators.n_layer_discriminator": [[108, "module-cyto_dl.nn.discriminators.n_layer_discriminator"]], "forward() (cyto_dl.nn.discriminators.n_layer_discriminator.nlayerdiscriminator method)": [[108, "cyto_dl.nn.discriminators.n_layer_discriminator.NLayerDiscriminator.forward"]], "cyto_dl.nn.head": [[109, "module-cyto_dl.nn.head"]], "basehead (class in cyto_dl.nn.head.base_head)": [[110, "cyto_dl.nn.head.base_head.BaseHead"]], "cyto_dl.nn.head.base_head": [[110, "module-cyto_dl.nn.head.base_head"]], "forward() (cyto_dl.nn.head.base_head.basehead method)": [[110, "cyto_dl.nn.head.base_head.BaseHead.forward"]], "generate_io_map() (cyto_dl.nn.head.base_head.basehead method)": [[110, "cyto_dl.nn.head.base_head.BaseHead.generate_io_map"]], "run_head() (cyto_dl.nn.head.base_head.basehead method)": [[110, "cyto_dl.nn.head.base_head.BaseHead.run_head"]], "update_params() (cyto_dl.nn.head.base_head.basehead method)": [[110, "cyto_dl.nn.head.base_head.BaseHead.update_params"]], "ganhead (class in cyto_dl.nn.head.gan_head)": [[111, "cyto_dl.nn.head.gan_head.GANHead"]], "cyto_dl.nn.head.gan_head": [[111, "module-cyto_dl.nn.head.gan_head"]], "forward() (cyto_dl.nn.head.gan_head.ganhead method)": [[111, "cyto_dl.nn.head.gan_head.GANHead.forward"]], "run_head() (cyto_dl.nn.head.gan_head.ganhead method)": [[111, "cyto_dl.nn.head.gan_head.GANHead.run_head"]], "ganhead_resize (class in cyto_dl.nn.head.gan_head_superres)": [[112, "cyto_dl.nn.head.gan_head_superres.GANHead_resize"]], "cyto_dl.nn.head.gan_head_superres": [[112, "module-cyto_dl.nn.head.gan_head_superres"]], "forward() (cyto_dl.nn.head.gan_head_superres.ganhead_resize method)": [[112, "cyto_dl.nn.head.gan_head_superres.GANHead_resize.forward"]], "maehead (class in cyto_dl.nn.head.mae_head)": [[113, "cyto_dl.nn.head.mae_head.MAEHead"]], "cyto_dl.nn.head.mae_head": [[113, "module-cyto_dl.nn.head.mae_head"]], "run_head() (cyto_dl.nn.head.mae_head.maehead method)": [[113, "cyto_dl.nn.head.mae_head.MAEHead.run_head"]], "maskhead (class in cyto_dl.nn.head.mask_head)": [[114, "cyto_dl.nn.head.mask_head.MaskHead"]], "cyto_dl.nn.head.mask_head": [[114, "module-cyto_dl.nn.head.mask_head"]], "run_head() (cyto_dl.nn.head.mask_head.maskhead method)": [[114, "cyto_dl.nn.head.mask_head.MaskHead.run_head"]], "resblockshead (class in cyto_dl.nn.head.res_blocks_head)": [[115, "cyto_dl.nn.head.res_blocks_head.ResBlocksHead"]], "cyto_dl.nn.head.res_blocks_head": [[115, "module-cyto_dl.nn.head.res_blocks_head"]], "forward() (cyto_dl.nn.head.res_blocks_head.resblockshead method)": [[115, "cyto_dl.nn.head.res_blocks_head.ResBlocksHead.forward"]], "vicreghead (class in cyto_dl.nn.head.vic_reg)": [[116, "cyto_dl.nn.head.vic_reg.VICRegHead"]], "cyto_dl.nn.head.vic_reg": [[116, "module-cyto_dl.nn.head.vic_reg"]], "run_head() (cyto_dl.nn.head.vic_reg.vicreghead method)": [[116, "cyto_dl.nn.head.vic_reg.VICRegHead.run_head"]], "hrskip (class in cyto_dl.nn.hr_skip)": [[117, "cyto_dl.nn.hr_skip.HRSkip"]], "cyto_dl.nn.hr_skip": [[117, "module-cyto_dl.nn.hr_skip"]], "forward() (cyto_dl.nn.hr_skip.hrskip method)": [[117, "cyto_dl.nn.hr_skip.HRSkip.forward"]], "cyto_dl.nn.losses": [[118, "module-cyto_dl.nn.losses"]], "adversarialloss (class in cyto_dl.nn.losses.adversarial_loss)": [[119, "cyto_dl.nn.losses.adversarial_loss.AdversarialLoss"]], "cyto_dl.nn.losses.adversarial_loss": [[119, "module-cyto_dl.nn.losses.adversarial_loss"]], "forward() (cyto_dl.nn.losses.adversarial_loss.adversarialloss method)": [[119, "cyto_dl.nn.losses.adversarial_loss.AdversarialLoss.forward"]], "chamferloss (class in cyto_dl.nn.losses.chamfer_loss)": [[120, "cyto_dl.nn.losses.chamfer_loss.ChamferLoss"]], "batch_pairwise_dist() (cyto_dl.nn.losses.chamfer_loss.chamferloss method)": [[120, "cyto_dl.nn.losses.chamfer_loss.ChamferLoss.batch_pairwise_dist"]], "batch_pairwise_dist2() (cyto_dl.nn.losses.chamfer_loss.chamferloss method)": [[120, "cyto_dl.nn.losses.chamfer_loss.ChamferLoss.batch_pairwise_dist2"]], "cyto_dl.nn.losses.chamfer_loss": [[120, "module-cyto_dl.nn.losses.chamfer_loss"]], "forward() (cyto_dl.nn.losses.chamfer_loss.chamferloss method)": [[120, "cyto_dl.nn.losses.chamfer_loss.ChamferLoss.forward"]], "cblogloss (class in cyto_dl.nn.losses.continuous_bernoulli)": [[121, "cyto_dl.nn.losses.continuous_bernoulli.CBLogLoss"]], "cyto_dl.nn.losses.continuous_bernoulli": [[121, "module-cyto_dl.nn.losses.continuous_bernoulli"]], "forward() (cyto_dl.nn.losses.continuous_bernoulli.cblogloss method)": [[121, "cyto_dl.nn.losses.continuous_bernoulli.CBLogLoss.forward"]], "cosineloss (class in cyto_dl.nn.losses.cosine_loss)": [[122, "cyto_dl.nn.losses.cosine_loss.CosineLoss"]], "cyto_dl.nn.losses.cosine_loss": [[122, "module-cyto_dl.nn.losses.cosine_loss"]], "forward() (cyto_dl.nn.losses.cosine_loss.cosineloss method)": [[122, "cyto_dl.nn.losses.cosine_loss.CosineLoss.forward"]], "ganloss (class in cyto_dl.nn.losses.gan_loss)": [[123, "cyto_dl.nn.losses.gan_loss.GANLoss"]], "pix2pixhd (class in cyto_dl.nn.losses.gan_loss)": [[123, "cyto_dl.nn.losses.gan_loss.Pix2PixHD"]], "cyto_dl.nn.losses.gan_loss": [[123, "module-cyto_dl.nn.losses.gan_loss"]], "get_feature_matching_loss() (cyto_dl.nn.losses.gan_loss.pix2pixhd method)": [[123, "cyto_dl.nn.losses.gan_loss.Pix2PixHD.get_feature_matching_loss"]], "get_gan_loss() (cyto_dl.nn.losses.gan_loss.pix2pixhd method)": [[123, "cyto_dl.nn.losses.gan_loss.Pix2PixHD.get_gan_loss"]], "get_target_tensor() (cyto_dl.nn.losses.gan_loss.ganloss method)": [[123, "cyto_dl.nn.losses.gan_loss.GANLoss.get_target_tensor"]], "gaussiannllloss (class in cyto_dl.nn.losses.gaussian_nll_loss)": [[124, "cyto_dl.nn.losses.gaussian_nll_loss.GaussianNLLLoss"]], "cyto_dl.nn.losses.gaussian_nll_loss": [[124, "module-cyto_dl.nn.losses.gaussian_nll_loss"]], "forward() (cyto_dl.nn.losses.gaussian_nll_loss.gaussiannllloss method)": [[124, "cyto_dl.nn.losses.gaussian_nll_loss.GaussianNLLLoss.forward"]], "cmap_loss (class in cyto_dl.nn.losses.loss_wrapper)": [[126, "cyto_dl.nn.losses.loss_wrapper.CMAP_loss"]], "losswrapper (class in cyto_dl.nn.losses.loss_wrapper)": [[126, "cyto_dl.nn.losses.loss_wrapper.LossWrapper"]], "cyto_dl.nn.losses.loss_wrapper": [[126, "module-cyto_dl.nn.losses.loss_wrapper"]], "thresholdloss (class in cyto_dl.nn.losses.threshold_loss)": [[128, "cyto_dl.nn.losses.threshold_loss.ThresholdLoss"]], "cyto_dl.nn.losses.threshold_loss": [[128, "module-cyto_dl.nn.losses.threshold_loss"]], "vicregloss (class in cyto_dl.nn.losses.vic_reg)": [[129, "cyto_dl.nn.losses.vic_reg.VICRegLoss"]], "cyto_dl.nn.losses.vic_reg": [[129, "module-cyto_dl.nn.losses.vic_reg"]], "forward() (cyto_dl.nn.losses.vic_reg.vicregloss method)": [[129, "cyto_dl.nn.losses.vic_reg.VICRegLoss.forward"]], "off_diagonal() (in module cyto_dl.nn.losses.vic_reg)": [[129, "cyto_dl.nn.losses.vic_reg.off_diagonal"]], "weibulllogloss (class in cyto_dl.nn.losses.weibull)": [[130, "cyto_dl.nn.losses.weibull.WeibullLogLoss"]], "cyto_dl.nn.losses.weibull": [[130, "module-cyto_dl.nn.losses.weibull"]], "forward() (cyto_dl.nn.losses.weibull.weibulllogloss method)": [[130, "cyto_dl.nn.losses.weibull.WeibullLogLoss.forward"]], "weibull_log_probs() (in module cyto_dl.nn.losses.weibull)": [[130, "cyto_dl.nn.losses.weibull.weibull_log_probs"]], "weightedmseloss (class in cyto_dl.nn.losses.weighted_mse_loss)": [[131, "cyto_dl.nn.losses.weighted_mse_loss.WeightedMSELoss"]], "cyto_dl.nn.losses.weighted_mse_loss": [[131, "module-cyto_dl.nn.losses.weighted_mse_loss"]], "forward() (cyto_dl.nn.losses.weighted_mse_loss.weightedmseloss method)": [[131, "cyto_dl.nn.losses.weighted_mse_loss.WeightedMSELoss.forward"]], "mlp (class in cyto_dl.nn.mlp)": [[132, "cyto_dl.nn.mlp.MLP"]], "cyto_dl.nn.mlp": [[132, "module-cyto_dl.nn.mlp"]], "forward() (cyto_dl.nn.mlp.mlp method)": [[132, "cyto_dl.nn.mlp.MLP.forward"]], "residualunit (class in cyto_dl.nn.res_unit)": [[138, "cyto_dl.nn.res_unit.ResidualUnit"]], "cyto_dl.nn.res_unit": [[138, "module-cyto_dl.nn.res_unit"]], "forward() (cyto_dl.nn.res_unit.residualunit method)": [[138, "cyto_dl.nn.res_unit.ResidualUnit.forward"]], "convpoolrelu (class in cyto_dl.nn.spatial_transformer)": [[139, "cyto_dl.nn.spatial_transformer.ConvPoolReLU"]], "stn (class in cyto_dl.nn.spatial_transformer)": [[139, "cyto_dl.nn.spatial_transformer.STN"]], "cyto_dl.nn.spatial_transformer": [[139, "module-cyto_dl.nn.spatial_transformer"]], "forward() (cyto_dl.nn.spatial_transformer.convpoolrelu method)": [[139, "cyto_dl.nn.spatial_transformer.ConvPoolReLU.forward"]], "forward() (cyto_dl.nn.spatial_transformer.stn method)": [[139, "cyto_dl.nn.spatial_transformer.STN.forward"]], "toggle() (cyto_dl.nn.spatial_transformer.stn method)": [[139, "cyto_dl.nn.spatial_transformer.STN.toggle"]], "trackclassifier (class in cyto_dl.nn.track_sequence_predictor)": [[140, "cyto_dl.nn.track_sequence_predictor.TrackClassifier"]], "cyto_dl.nn.track_sequence_predictor": [[140, "module-cyto_dl.nn.track_sequence_predictor"]], "forward() (cyto_dl.nn.track_sequence_predictor.trackclassifier method)": [[140, "cyto_dl.nn.track_sequence_predictor.TrackClassifier.forward"]], "positionalencoding1d() (in module cyto_dl.nn.track_sequence_predictor)": [[140, "cyto_dl.nn.track_sequence_predictor.positionalencoding1d"]], "cyto_dl.nn.vits": [[141, "module-cyto_dl.nn.vits"]], "cyto_dl.nn.vits.blocks": [[142, "module-cyto_dl.nn.vits.blocks"]], "crossattention (class in cyto_dl.nn.vits.blocks.cross_attention)": [[143, "cyto_dl.nn.vits.blocks.cross_attention.CrossAttention"]], "crossattentionblock (class in cyto_dl.nn.vits.blocks.cross_attention)": [[143, "cyto_dl.nn.vits.blocks.cross_attention.CrossAttentionBlock"]], "crossselfblock (class in cyto_dl.nn.vits.blocks.cross_attention)": [[143, "cyto_dl.nn.vits.blocks.cross_attention.CrossSelfBlock"]], "mlp (class in cyto_dl.nn.vits.blocks.cross_attention)": [[143, "cyto_dl.nn.vits.blocks.cross_attention.Mlp"]], "cyto_dl.nn.vits.blocks.cross_attention": [[143, "module-cyto_dl.nn.vits.blocks.cross_attention"]], "forward() (cyto_dl.nn.vits.blocks.cross_attention.crossattention method)": [[143, "cyto_dl.nn.vits.blocks.cross_attention.CrossAttention.forward"]], "forward() (cyto_dl.nn.vits.blocks.cross_attention.crossattentionblock method)": [[143, "cyto_dl.nn.vits.blocks.cross_attention.CrossAttentionBlock.forward"]], "forward() (cyto_dl.nn.vits.blocks.cross_attention.crossselfblock method)": [[143, "cyto_dl.nn.vits.blocks.cross_attention.CrossSelfBlock.forward"]], "forward() (cyto_dl.nn.vits.blocks.cross_attention.mlp method)": [[143, "cyto_dl.nn.vits.blocks.cross_attention.Mlp.forward"]], "intermediateweigher (class in cyto_dl.nn.vits.blocks.intermediate_weigher)": [[144, "cyto_dl.nn.vits.blocks.intermediate_weigher.IntermediateWeigher"]], "cyto_dl.nn.vits.blocks.intermediate_weigher": [[144, "module-cyto_dl.nn.vits.blocks.intermediate_weigher"]], "forward() (cyto_dl.nn.vits.blocks.intermediate_weigher.intermediateweigher method)": [[144, "cyto_dl.nn.vits.blocks.intermediate_weigher.IntermediateWeigher.forward"]], "hierablock (class in cyto_dl.nn.vits.blocks.masked_unit_attention)": [[145, "cyto_dl.nn.vits.blocks.masked_unit_attention.HieraBlock"]], "maskunitattention (class in cyto_dl.nn.vits.blocks.masked_unit_attention)": [[145, "cyto_dl.nn.vits.blocks.masked_unit_attention.MaskUnitAttention"]], "cyto_dl.nn.vits.blocks.masked_unit_attention": [[145, "module-cyto_dl.nn.vits.blocks.masked_unit_attention"]], "forward() (cyto_dl.nn.vits.blocks.masked_unit_attention.hierablock method)": [[145, "cyto_dl.nn.vits.blocks.masked_unit_attention.HieraBlock.forward"]], "forward() (cyto_dl.nn.vits.blocks.masked_unit_attention.maskunitattention method)": [[145, "cyto_dl.nn.vits.blocks.masked_unit_attention.MaskUnitAttention.forward"]], "cyto_dl.nn.vits.blocks.patchify": [[146, "module-cyto_dl.nn.vits.blocks.patchify"]], "patchify (class in cyto_dl.nn.vits.blocks.patchify.patchify)": [[147, "cyto_dl.nn.vits.blocks.patchify.patchify.Patchify"]], "create_img2token() (cyto_dl.nn.vits.blocks.patchify.patchify.patchify method)": [[147, "cyto_dl.nn.vits.blocks.patchify.patchify.Patchify.create_img2token"]], "cyto_dl.nn.vits.blocks.patchify.patchify": [[147, "module-cyto_dl.nn.vits.blocks.patchify.patchify"]], "extract_visible_tokens() (cyto_dl.nn.vits.blocks.patchify.patchify.patchify method)": [[147, "cyto_dl.nn.vits.blocks.patchify.patchify.Patchify.extract_visible_tokens"]], "get_mask_args() (cyto_dl.nn.vits.blocks.patchify.patchify.patchify method)": [[147, "cyto_dl.nn.vits.blocks.patchify.patchify.Patchify.get_mask_args"]], "img2token (cyto_dl.nn.vits.blocks.patchify.patchify.patchify property)": [[147, "cyto_dl.nn.vits.blocks.patchify.patchify.Patchify.img2token"]], "patchifybase (class in cyto_dl.nn.vits.blocks.patchify.patchify_base)": [[148, "cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase"]], "create_conv() (cyto_dl.nn.vits.blocks.patchify.patchify_base.patchifybase method)": [[148, "cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase.create_conv"]], "create_patch2img() (cyto_dl.nn.vits.blocks.patchify.patchify_base.patchifybase method)": [[148, "cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase.create_patch2img"]], "cyto_dl.nn.vits.blocks.patchify.patchify_base": [[148, "module-cyto_dl.nn.vits.blocks.patchify.patchify_base"]], "extract_visible_tokens() (cyto_dl.nn.vits.blocks.patchify.patchify_base.patchifybase method)": [[148, "cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase.extract_visible_tokens"]], "forward() (cyto_dl.nn.vits.blocks.patchify.patchify_base.patchifybase method)": [[148, "cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase.forward"]], "get_mask() (cyto_dl.nn.vits.blocks.patchify.patchify_base.patchifybase method)": [[148, "cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase.get_mask"]], "get_mask_args() (cyto_dl.nn.vits.blocks.patchify.patchify_base.patchifybase method)": [[148, "cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase.get_mask_args"]], "img2token (cyto_dl.nn.vits.blocks.patchify.patchify_base.patchifybase property)": [[148, "cyto_dl.nn.vits.blocks.patchify.patchify_base.PatchifyBase.img2token"]], "patchifyhiera (class in cyto_dl.nn.vits.blocks.patchify.patchify_hiera)": [[149, "cyto_dl.nn.vits.blocks.patchify.patchify_hiera.PatchifyHiera"]], "create_img2token() (cyto_dl.nn.vits.blocks.patchify.patchify_hiera.patchifyhiera method)": [[149, "cyto_dl.nn.vits.blocks.patchify.patchify_hiera.PatchifyHiera.create_img2token"]], "cyto_dl.nn.vits.blocks.patchify.patchify_hiera": [[149, "module-cyto_dl.nn.vits.blocks.patchify.patchify_hiera"]], "extract_visible_tokens() (cyto_dl.nn.vits.blocks.patchify.patchify_hiera.patchifyhiera method)": [[149, "cyto_dl.nn.vits.blocks.patchify.patchify_hiera.PatchifyHiera.extract_visible_tokens"]], "get_mask_args() (cyto_dl.nn.vits.blocks.patchify.patchify_hiera.patchifyhiera method)": [[149, "cyto_dl.nn.vits.blocks.patchify.patchify_hiera.PatchifyHiera.get_mask_args"]], "img2token (cyto_dl.nn.vits.blocks.patchify.patchify_hiera.patchifyhiera property)": [[149, "cyto_dl.nn.vits.blocks.patchify.patchify_hiera.PatchifyHiera.img2token"]], "take_indexes_mask() (in module cyto_dl.nn.vits.blocks.patchify.patchify_hiera)": [[149, "cyto_dl.nn.vits.blocks.patchify.patchify_hiera.take_indexes_mask"]], "crossmae_decoder (class in cyto_dl.nn.vits.decoder)": [[150, "cyto_dl.nn.vits.decoder.CrossMAE_Decoder"]], "mae_decoder (class in cyto_dl.nn.vits.decoder)": [[150, "cyto_dl.nn.vits.decoder.MAE_Decoder"]], "add_mask_tokens() (cyto_dl.nn.vits.decoder.mae_decoder method)": [[150, "cyto_dl.nn.vits.decoder.MAE_Decoder.add_mask_tokens"]], "adjust_indices_for_cls() (cyto_dl.nn.vits.decoder.mae_decoder method)": [[150, "cyto_dl.nn.vits.decoder.MAE_Decoder.adjust_indices_for_cls"]], "cyto_dl.nn.vits.decoder": [[150, "module-cyto_dl.nn.vits.decoder"]], "forward() (cyto_dl.nn.vits.decoder.crossmae_decoder method)": [[150, "cyto_dl.nn.vits.decoder.CrossMAE_Decoder.forward"]], "forward() (cyto_dl.nn.vits.decoder.mae_decoder method)": [[150, "cyto_dl.nn.vits.decoder.MAE_Decoder.forward"]], "init_weight() (cyto_dl.nn.vits.decoder.mae_decoder method)": [[150, "cyto_dl.nn.vits.decoder.MAE_Decoder.init_weight"]], "hieraencoder (class in cyto_dl.nn.vits.encoder)": [[151, "cyto_dl.nn.vits.encoder.HieraEncoder"]], "jepaencoder (class in cyto_dl.nn.vits.encoder)": [[151, "cyto_dl.nn.vits.encoder.JEPAEncoder"]], "mae_encoder (class in cyto_dl.nn.vits.encoder)": [[151, "cyto_dl.nn.vits.encoder.MAE_Encoder"]], "spatialmerger (class in cyto_dl.nn.vits.encoder)": [[151, "cyto_dl.nn.vits.encoder.SpatialMerger"]], "cyto_dl.nn.vits.encoder": [[151, "module-cyto_dl.nn.vits.encoder"]], "forward() (cyto_dl.nn.vits.encoder.hieraencoder method)": [[151, "cyto_dl.nn.vits.encoder.HieraEncoder.forward"]], "forward() (cyto_dl.nn.vits.encoder.jepaencoder method)": [[151, "cyto_dl.nn.vits.encoder.JEPAEncoder.forward"]], "forward() (cyto_dl.nn.vits.encoder.mae_encoder method)": [[151, "cyto_dl.nn.vits.encoder.MAE_Encoder.forward"]], "forward() (cyto_dl.nn.vits.encoder.spatialmerger method)": [[151, "cyto_dl.nn.vits.encoder.SpatialMerger.forward"]], "init_weight() (cyto_dl.nn.vits.encoder.mae_encoder method)": [[151, "cyto_dl.nn.vits.encoder.MAE_Encoder.init_weight"]], "hieramae (class in cyto_dl.nn.vits.mae)": [[152, "cyto_dl.nn.vits.mae.HieraMAE"]], "mae (class in cyto_dl.nn.vits.mae)": [[152, "cyto_dl.nn.vits.mae.MAE"]], "mae_base (class in cyto_dl.nn.vits.mae)": [[152, "cyto_dl.nn.vits.mae.MAE_Base"]], "cyto_dl.nn.vits.mae": [[152, "module-cyto_dl.nn.vits.mae"]], "decoder (cyto_dl.nn.vits.mae.hieramae property)": [[152, "cyto_dl.nn.vits.mae.HieraMAE.decoder"]], "decoder (cyto_dl.nn.vits.mae.mae property)": [[152, "cyto_dl.nn.vits.mae.MAE.decoder"]], "decoder (cyto_dl.nn.vits.mae.mae_base property)": [[152, "cyto_dl.nn.vits.mae.MAE_Base.decoder"]], "encoder (cyto_dl.nn.vits.mae.hieramae property)": [[152, "cyto_dl.nn.vits.mae.HieraMAE.encoder"]], "encoder (cyto_dl.nn.vits.mae.mae property)": [[152, "cyto_dl.nn.vits.mae.MAE.encoder"]], "encoder (cyto_dl.nn.vits.mae.mae_base property)": [[152, "cyto_dl.nn.vits.mae.MAE_Base.encoder"]], "forward() (cyto_dl.nn.vits.mae.mae_base method)": [[152, "cyto_dl.nn.vits.mae.MAE_Base.forward"]], "init_decoder() (cyto_dl.nn.vits.mae.mae_base method)": [[152, "cyto_dl.nn.vits.mae.MAE_Base.init_decoder"]], "init_encoder() (cyto_dl.nn.vits.mae.mae_base method)": [[152, "cyto_dl.nn.vits.mae.MAE_Base.init_encoder"]], "iwmpredictor (class in cyto_dl.nn.vits.predictor)": [[153, "cyto_dl.nn.vits.predictor.IWMPredictor"]], "jepapredictor (class in cyto_dl.nn.vits.predictor)": [[153, "cyto_dl.nn.vits.predictor.JEPAPredictor"]], "cyto_dl.nn.vits.predictor": [[153, "module-cyto_dl.nn.vits.predictor"]], "forward() (cyto_dl.nn.vits.predictor.iwmpredictor method)": [[153, "cyto_dl.nn.vits.predictor.IWMPredictor.forward"]], "forward() (cyto_dl.nn.vits.predictor.jepapredictor method)": [[153, "cyto_dl.nn.vits.predictor.JEPAPredictor.forward"]], "init_weight() (cyto_dl.nn.vits.predictor.jepapredictor method)": [[153, "cyto_dl.nn.vits.predictor.JEPAPredictor.init_weight"]], "predict_target_features() (cyto_dl.nn.vits.predictor.jepapredictor method)": [[153, "cyto_dl.nn.vits.predictor.JEPAPredictor.predict_target_features"]], "encodedskip (class in cyto_dl.nn.vits.seg)": [[154, "cyto_dl.nn.vits.seg.EncodedSkip"]], "seg_vit (class in cyto_dl.nn.vits.seg)": [[154, "cyto_dl.nn.vits.seg.Seg_ViT"]], "superresdecoder (class in cyto_dl.nn.vits.seg)": [[154, "cyto_dl.nn.vits.seg.SuperresDecoder"]], "cyto_dl.nn.vits.seg": [[154, "module-cyto_dl.nn.vits.seg"]], "forward() (cyto_dl.nn.vits.seg.encodedskip method)": [[154, "cyto_dl.nn.vits.seg.EncodedSkip.forward"]], "forward() (cyto_dl.nn.vits.seg.seg_vit method)": [[154, "cyto_dl.nn.vits.seg.Seg_ViT.forward"]], "forward() (cyto_dl.nn.vits.seg.superresdecoder method)": [[154, "cyto_dl.nn.vits.seg.SuperresDecoder.forward"]], "cyto_dl.nn.vits.utils": [[155, "module-cyto_dl.nn.vits.utils"]], "get_positional_embedding() (in module cyto_dl.nn.vits.utils)": [[155, "cyto_dl.nn.vits.utils.get_positional_embedding"]], "match_tuple_dimensions() (in module cyto_dl.nn.vits.utils)": [[155, "cyto_dl.nn.vits.utils.match_tuple_dimensions"]], "random_indexes() (in module cyto_dl.nn.vits.utils)": [[155, "cyto_dl.nn.vits.utils.random_indexes"]], "take_indexes() (in module cyto_dl.nn.vits.utils)": [[155, "cyto_dl.nn.vits.utils.take_indexes"]], "cyto_dl.point_cloud": [[156, "module-cyto_dl.point_cloud"]], "cyto_dl.train": [[159, "module-cyto_dl.train"]], "main() (in module cyto_dl.train)": [[159, "cyto_dl.train.main"]], "cyto_dl.utils": [[160, "module-cyto_dl.utils"]], "create_dataloader() (in module cyto_dl.utils.array)": [[161, "cyto_dl.utils.array.create_dataloader"]], "cyto_dl.utils.array": [[161, "module-cyto_dl.utils.array"]], "extract_array_predictions() (in module cyto_dl.utils.array)": [[161, "cyto_dl.utils.array.extract_array_predictions"]], "cyto_dl.utils.checkpoint": [[162, "module-cyto_dl.utils.checkpoint"]], "load_checkpoint() (in module cyto_dl.utils.checkpoint)": [[162, "cyto_dl.utils.checkpoint.load_checkpoint"]], "cyto_dl.utils.config": [[163, "module-cyto_dl.utils.config"]], "is_config() (in module cyto_dl.utils.config)": [[163, "cyto_dl.utils.config.is_config"]], "kv_to_dict() (in module cyto_dl.utils.config)": [[163, "cyto_dl.utils.config.kv_to_dict"]], "remove_aux_key() (in module cyto_dl.utils.config)": [[163, "cyto_dl.utils.config.remove_aux_key"]], "cyto_dl.utils.download_test_data": [[164, "module-cyto_dl.utils.download_test_data"]], "delete_test_data() (in module cyto_dl.utils.download_test_data)": [[164, "cyto_dl.utils.download_test_data.delete_test_data"]], "download_test_data() (in module cyto_dl.utils.download_test_data)": [[164, "cyto_dl.utils.download_test_data.download_test_data"]], "parse_s3_path() (in module cyto_dl.utils.download_test_data)": [[164, "cyto_dl.utils.download_test_data.parse_s3_path"]], "setup_paths() (in module cyto_dl.utils.download_test_data)": [[164, "cyto_dl.utils.download_test_data.setup_paths"]], "dummydatamodule (class in cyto_dl.utils.dummy_dataset)": [[165, "cyto_dl.utils.dummy_dataset.DummyDatamodule"]], "dummydataset (class in cyto_dl.utils.dummy_dataset)": [[165, "cyto_dl.utils.dummy_dataset.DummyDataset"]], "cyto_dl.utils.dummy_dataset": [[165, "module-cyto_dl.utils.dummy_dataset"]], "generate_img() (cyto_dl.utils.dummy_dataset.dummydataset method)": [[165, "cyto_dl.utils.dummy_dataset.DummyDataset.generate_img"]], "get_dataloader() (cyto_dl.utils.dummy_dataset.dummydatamodule method)": [[165, "cyto_dl.utils.dummy_dataset.DummyDatamodule.get_dataloader"]], "predict_dataloader() (cyto_dl.utils.dummy_dataset.dummydatamodule method)": [[165, "cyto_dl.utils.dummy_dataset.DummyDatamodule.predict_dataloader"]], "test_dataloader() (cyto_dl.utils.dummy_dataset.dummydatamodule method)": [[165, "cyto_dl.utils.dummy_dataset.DummyDatamodule.test_dataloader"]], "train_dataloader() (cyto_dl.utils.dummy_dataset.dummydatamodule method)": [[165, "cyto_dl.utils.dummy_dataset.DummyDatamodule.train_dataloader"]], "val_dataloader() (cyto_dl.utils.dummy_dataset.dummydatamodule method)": [[165, "cyto_dl.utils.dummy_dataset.DummyDatamodule.val_dataloader"]], "cyto_dl.utils.pylogger": [[166, "module-cyto_dl.utils.pylogger"]], "get_pylogger() (in module cyto_dl.utils.pylogger)": [[166, "cyto_dl.utils.pylogger.get_pylogger"]], "cyto_dl.utils.rich_utils": [[167, "module-cyto_dl.utils.rich_utils"]], "enforce_tags() (in module cyto_dl.utils.rich_utils)": [[167, "cyto_dl.utils.rich_utils.enforce_tags"]], "print_config_tree() (in module cyto_dl.utils.rich_utils)": [[167, "cyto_dl.utils.rich_utils.print_config_tree"]], "rotationmodule (class in cyto_dl.utils.rotation)": [[168, "cyto_dl.utils.rotation.RotationModule"]], "compute_rotation_matrix() (cyto_dl.utils.rotation.rotationmodule method)": [[168, "cyto_dl.utils.rotation.RotationModule.compute_rotation_matrix"]], "cyto_dl.utils.rotation": [[168, "module-cyto_dl.utils.rotation"]], "get_rotation_matrix_so2() (in module cyto_dl.utils.rotation)": [[168, "cyto_dl.utils.rotation.get_rotation_matrix_so2"]], "get_rotation_matrix_so3() (in module cyto_dl.utils.rotation)": [[168, "cyto_dl.utils.rotation.get_rotation_matrix_so3"]], "close_loggers() (in module cyto_dl.utils.template_utils)": [[172, "cyto_dl.utils.template_utils.close_loggers"]], "cyto_dl.utils.template_utils": [[172, "module-cyto_dl.utils.template_utils"]], "extras() (in module cyto_dl.utils.template_utils)": [[172, "cyto_dl.utils.template_utils.extras"]], "get_metric_value() (in module cyto_dl.utils.template_utils)": [[172, "cyto_dl.utils.template_utils.get_metric_value"]], "instantiate_callbacks() (in module cyto_dl.utils.template_utils)": [[172, "cyto_dl.utils.template_utils.instantiate_callbacks"]], "instantiate_loggers() (in module cyto_dl.utils.template_utils)": [[172, "cyto_dl.utils.template_utils.instantiate_loggers"]], "log_hyperparameters() (in module cyto_dl.utils.template_utils)": [[172, "cyto_dl.utils.template_utils.log_hyperparameters"]], "save_file() (in module cyto_dl.utils.template_utils)": [[172, "cyto_dl.utils.template_utils.save_file"]], "task_wrapper() (in module cyto_dl.utils.template_utils)": [[172, "cyto_dl.utils.template_utils.task_wrapper"]]}})
\ No newline at end of file
diff --git a/using_examples.html b/using_examples.html
index d745f9edc..3f65eb713 100644
--- a/using_examples.html
+++ b/using_examples.html
@@ -350,14 +350,21 @@
 </li>
 <li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.nn.vits.html">cyto_dl.nn.vits package</a><input class="toctree-checkbox" id="toctree-checkbox-31" name="toctree-checkbox-31" role="switch" type="checkbox"/><label for="toctree-checkbox-31"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.html">cyto_dl.nn.vits.blocks package</a><input class="toctree-checkbox" id="toctree-checkbox-32" name="toctree-checkbox-32" role="switch" type="checkbox"/><label for="toctree-checkbox-32"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l6 has-children"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify.html">cyto_dl.nn.vits.blocks.patchify.patchify module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_base.html">cyto_dl.nn.vits.blocks.patchify.patchify_base module</a></li>
+<li class="toctree-l7"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.patchify_hiera.html">cyto_dl.nn.vits.blocks.patchify.patchify_hiera module</a></li>
+</ul>
+</li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.cross_attention.html">cyto_dl.nn.vits.blocks.cross_attention module</a></li>
 <li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.intermediate_weigher.html">cyto_dl.nn.vits.blocks.intermediate_weigher module</a></li>
-<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.patchify.html">cyto_dl.nn.vits.blocks.patchify module</a></li>
+<li class="toctree-l6"><a class="reference internal" href="cyto_dl.nn.vits.blocks.masked_unit_attention.html">cyto_dl.nn.vits.blocks.masked_unit_attention module</a></li>
 </ul>
 </li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.cross_mae.html">cyto_dl.nn.vits.cross_mae module</a></li>
-<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.jepa.html">cyto_dl.nn.vits.jepa module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.decoder.html">cyto_dl.nn.vits.decoder module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.encoder.html">cyto_dl.nn.vits.encoder module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.mae.html">cyto_dl.nn.vits.mae module</a></li>
+<li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.predictor.html">cyto_dl.nn.vits.predictor module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.seg.html">cyto_dl.nn.vits.seg module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.nn.vits.utils.html">cyto_dl.nn.vits.utils module</a></li>
 </ul>
@@ -369,15 +376,15 @@
 <li class="toctree-l4"><a class="reference internal" href="cyto_dl.nn.track_sequence_predictor.html">cyto_dl.nn.track_sequence_predictor module</a></li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-33" name="toctree-checkbox-33" role="switch" type="checkbox"/><label for="toctree-checkbox-33"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.point_cloud.html">cyto_dl.point_cloud package</a><input class="toctree-checkbox" id="toctree-checkbox-34" name="toctree-checkbox-34" role="switch" type="checkbox"/><label for="toctree-checkbox-34"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.point_cloud.io.html">cyto_dl.point_cloud.io package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.point_cloud.io.read_pcloud.html">cyto_dl.point_cloud.io.read_pcloud module</a></li>
 </ul>
 </li>
 </ul>
 </li>
-<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-35" name="toctree-checkbox-35" role="switch" type="checkbox"/><label for="toctree-checkbox-35"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
-<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l3 has-children"><a class="reference internal" href="cyto_dl.utils.html">cyto_dl.utils package</a><input class="toctree-checkbox" id="toctree-checkbox-36" name="toctree-checkbox-36" role="switch" type="checkbox"/><label for="toctree-checkbox-36"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
+<li class="toctree-l4 has-children"><a class="reference internal" href="cyto_dl.utils.spharm.html">cyto_dl.utils.spharm package</a><input class="toctree-checkbox" id="toctree-checkbox-37" name="toctree-checkbox-37" role="switch" type="checkbox"/><label for="toctree-checkbox-37"><div class="visually-hidden">Toggle child pages in navigation</div><i class="icon"><svg><use href="#svg-arrow-right"></use></svg></i></label><ul>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.reconstruction.html">cyto_dl.utils.spharm.reconstruction module</a></li>
 <li class="toctree-l5"><a class="reference internal" href="cyto_dl.utils.spharm.rotation.html">cyto_dl.utils.spharm.rotation module</a></li>
 </ul>