monai-weekly 1.5.dev2522__py3-none-any.whl → 1.6.dev2524__py3-none-any.whl

monai/__init__.py

@@ -136,4 +136,4 @@ except BaseException:
 
     if MONAIEnvVars.debug():
         raise
-__commit_id__ = "c3a317d2bcb486199f40bda0d722a41e3869712a"
+__commit_id__ = "d388d1c6fec8cb3a0eebee5b5a0b9776ca59ca83"

monai/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2025-06-01T02:46:06+0000",
+ "date": "2025-06-15T02:40:48+0000",
  "dirty": false,
  "error": null,
- "full-revisionid": "3e75ca0d8f4cf7ce1e4b3e01ef6058c56af87fdd",
- "version": "1.5.dev2522"
+ "full-revisionid": "1be81d981f92f8fa2fb1c27dfe5dad8c2e371624",
+ "version": "1.6.dev2524"
 }
 '''  # END VERSION_JSON
 

monai/inferers/merger.py

@@ -21,7 +21,14 @@ from typing import TYPE_CHECKING, Any
 import numpy as np
 import torch
 
-from monai.utils import ensure_tuple_size, get_package_version, optional_import, require_pkg, version_geq
+from monai.utils import (
+    deprecated_arg,
+    ensure_tuple_size,
+    get_package_version,
+    optional_import,
+    require_pkg,
+    version_geq,
+)
 
 if TYPE_CHECKING:
     import zarr
@@ -218,15 +225,41 @@ class ZarrAvgMerger(Merger):
         store: the zarr store to save the final results. Default is "merged.zarr".
         value_store: the zarr store to save the value aggregating tensor. Default is a temporary store.
         count_store: the zarr store to save the sample counting tensor. Default is a temporary store.
-        compressor: the compressor for final merged zarr array. Default is "default".
+        compressor: the compressor for final merged zarr array. Default is None.
+            Deprecated since 1.5.0 and will be removed in 1.7.0. Use codecs instead.
         value_compressor: the compressor for value aggregating zarr array. Default is None.
+            Deprecated since 1.5.0 and will be removed in 1.7.0. Use value_codecs instead.
         count_compressor: the compressor for sample counting zarr array. Default is None.
+            Deprecated since 1.5.0 and will be removed in 1.7.0. Use count_codecs instead.
+        codecs: the codecs for final merged zarr array. Default is None.
+            For zarr v3, this is a list of codec configurations. See zarr documentation for details.
+        value_codecs: the codecs for value aggregating zarr array. Default is None.
+            For zarr v3, this is a list of codec configurations. See zarr documentation for details.
+        count_codecs: the codecs for sample counting zarr array. Default is None.
+            For zarr v3, this is a list of codec configurations. See zarr documentation for details.
         chunks : int or tuple of ints that defines the chunk shape, or boolean. Default is True.
             If True, chunk shape will be guessed from `shape` and `dtype`.
             If False, it will be set to `shape`, i.e., single chunk for the whole array.
             If an int, the chunk size in each dimension will be given by the value of `chunks`.
     """
 
+    @deprecated_arg(
+        name="compressor", since="1.5.0", removed="1.7.0", new_name="codecs", msg_suffix="Please use 'codecs' instead."
+    )
+    @deprecated_arg(
+        name="value_compressor",
+        since="1.5.0",
+        removed="1.7.0",
+        new_name="value_codecs",
+        msg_suffix="Please use 'value_codecs' instead.",
+    )
+    @deprecated_arg(
+        name="count_compressor",
+        since="1.5.0",
+        removed="1.7.0",
+        new_name="count_codecs",
+        msg_suffix="Please use 'count_codecs' instead.",
+    )
     def __init__(
         self,
         merged_shape: Sequence[int],
@@ -240,6 +273,9 @@ class ZarrAvgMerger(Merger):
         compressor: str | None = None,
         value_compressor: str | None = None,
         count_compressor: str | None = None,
+        codecs: list | None = None,
+        value_codecs: list | None = None,
+        count_codecs: list | None = None,
         chunks: Sequence[int] | bool = True,
         thread_locking: bool = True,
     ) -> None:
@@ -251,7 +287,11 @@ class ZarrAvgMerger(Merger):
         self.count_dtype = count_dtype
         self.store = store
         self.tmpdir: TemporaryDirectory | None
-        if version_geq(get_package_version("zarr"), "3.0.0"):
+
+        # Handle zarr v3 vs older versions
+        is_zarr_v3 = version_geq(get_package_version("zarr"), "3.0.0")
+
+        if is_zarr_v3:
             if value_store is None:
                 self.tmpdir = TemporaryDirectory()
                 self.value_store = zarr.storage.LocalStore(self.tmpdir.name)  # type: ignore
@@ -266,34 +306,119 @@ class ZarrAvgMerger(Merger):
             self.tmpdir = None
             self.value_store = zarr.storage.TempStore() if value_store is None else value_store  # type: ignore
             self.count_store = zarr.storage.TempStore() if count_store is None else count_store  # type: ignore
+
         self.chunks = chunks
-        self.compressor = compressor
-        self.value_compressor = value_compressor
-        self.count_compressor = count_compressor
-        self.output = zarr.empty(
-            shape=self.merged_shape,
-            chunks=self.chunks,
-            dtype=self.output_dtype,
-            compressor=self.compressor,
-            store=self.store,
-            overwrite=True,
-        )
-        self.values = zarr.zeros(
-            shape=self.merged_shape,
-            chunks=self.chunks,
-            dtype=self.value_dtype,
-            compressor=self.value_compressor,
-            store=self.value_store,
-            overwrite=True,
-        )
-        self.counts = zarr.zeros(
-            shape=self.merged_shape,
-            chunks=self.chunks,
-            dtype=self.count_dtype,
-            compressor=self.count_compressor,
-            store=self.count_store,
-            overwrite=True,
-        )
+
+        # Handle compressor/codecs based on zarr version
+        is_zarr_v3 = version_geq(get_package_version("zarr"), "3.0.0")
+
+        # Initialize codecs/compressor attributes with proper types
+        self.codecs: list | None = None
+        self.value_codecs: list | None = None
+        self.count_codecs: list | None = None
+
+        if is_zarr_v3:
+            # For zarr v3, use codecs or convert compressor to codecs
+            if codecs is not None:
+                self.codecs = codecs
+            elif compressor is not None:
+                # Convert compressor to codec format
+                if isinstance(compressor, (list, tuple)):
+                    self.codecs = compressor
+                else:
+                    self.codecs = [compressor]
+            else:
+                self.codecs = None
+
+            if value_codecs is not None:
+                self.value_codecs = value_codecs
+            elif value_compressor is not None:
+                if isinstance(value_compressor, (list, tuple)):
+                    self.value_codecs = value_compressor
+                else:
+                    self.value_codecs = [value_compressor]
+            else:
+                self.value_codecs = None
+
+            if count_codecs is not None:
+                self.count_codecs = count_codecs
+            elif count_compressor is not None:
+                if isinstance(count_compressor, (list, tuple)):
+                    self.count_codecs = count_compressor
+                else:
+                    self.count_codecs = [count_compressor]
+            else:
+                self.count_codecs = None
+        else:
+            # For zarr v2, use compressors
+            if codecs is not None:
+                # If codecs are specified in v2, use the first codec as compressor
+                self.codecs = codecs[0] if isinstance(codecs, (list, tuple)) else codecs
+            else:
+                self.codecs = compressor  # type: ignore[assignment]
+
+            if value_codecs is not None:
+                self.value_codecs = value_codecs[0] if isinstance(value_codecs, (list, tuple)) else value_codecs
+            else:
+                self.value_codecs = value_compressor  # type: ignore[assignment]
+
+            if count_codecs is not None:
+                self.count_codecs = count_codecs[0] if isinstance(count_codecs, (list, tuple)) else count_codecs
+            else:
+                self.count_codecs = count_compressor  # type: ignore[assignment]
+
+        # Create zarr arrays with appropriate parameters based on version
+        if is_zarr_v3:
+            self.output = zarr.empty(
+                shape=self.merged_shape,
+                chunks=self.chunks,
+                dtype=self.output_dtype,
+                codecs=self.codecs,
+                store=self.store,
+                overwrite=True,
+            )
+            self.values = zarr.zeros(
+                shape=self.merged_shape,
+                chunks=self.chunks,
+                dtype=self.value_dtype,
+                codecs=self.value_codecs,
+                store=self.value_store,
+                overwrite=True,
+            )
+            self.counts = zarr.zeros(
+                shape=self.merged_shape,
+                chunks=self.chunks,
+                dtype=self.count_dtype,
+                codecs=self.count_codecs,
+                store=self.count_store,
+                overwrite=True,
+            )
+        else:
+            self.output = zarr.empty(
+                shape=self.merged_shape,
+                chunks=self.chunks,
+                dtype=self.output_dtype,
+                compressor=self.codecs,
+                store=self.store,
+                overwrite=True,
+            )
+            self.values = zarr.zeros(
+                shape=self.merged_shape,
+                chunks=self.chunks,
+                dtype=self.value_dtype,
+                compressor=self.value_codecs,
+                store=self.value_store,
+                overwrite=True,
+            )
+            self.counts = zarr.zeros(
+                shape=self.merged_shape,
+                chunks=self.chunks,
+                dtype=self.count_dtype,
+                compressor=self.count_codecs,
+                store=self.count_store,
+                overwrite=True,
+            )
+
         self.lock: threading.Lock | nullcontext
         if thread_locking:
             # use lock to protect the in-place addition during aggregation

{monai_weekly-1.5.dev2522.dist-info → monai_weekly-1.6.dev2524.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: monai-weekly
-Version: 1.5.dev2522
+Version: 1.6.dev2524
 Summary: AI Toolkit for Healthcare Imaging
 Home-page: https://monai.io/
 Author: MONAI Consortium
@@ -215,7 +215,7 @@ Please refer to [the installation guide](https://docs.monai.io/en/latest/install
 
 ## Getting Started
 
-[MedNIST demo](https://colab.research.google.com/drive/1wy8XUSnNWlhDNazFdvGBHLfdkGvOHBKe) and [MONAI for PyTorch Users](https://colab.research.google.com/drive/1boqy7ENpKrqaJoxFlbHIBnIODAs1Ih1T) are available on Colab.
+[MedNIST demo](https://colab.research.google.com/github/Project-MONAI/tutorials/blob/main/2d_classification/mednist_tutorial.ipynb) and [MONAI for PyTorch Users](https://colab.research.google.com/github/Project-MONAI/tutorials/blob/main/modules/developer_guide.ipynb) are available on Colab.
 
 Examples and notebook tutorials are located at [Project-MONAI/tutorials](https://github.com/Project-MONAI/tutorials).
 
@@ -236,7 +236,7 @@ For guidance on making a contribution to MONAI, see the [contributing guidelines
 
 ## Community
 
-Join the conversation on Twitter/X [@ProjectMONAI](https://twitter.com/ProjectMONAI) or join our [Slack channel](https://forms.gle/QTxJq3hFictp31UM9).
+Join the conversation on Twitter/X [@ProjectMONAI](https://twitter.com/ProjectMONAI), [LinkedIn](https://www.linkedin.com/company/projectmonai), or join our [Slack channel](https://forms.gle/QTxJq3hFictp31UM9).
 
 Ask and answer questions over on [MONAI's GitHub Discussions tab](https://github.com/Project-MONAI/MONAI/discussions).
 

{monai_weekly-1.5.dev2522.dist-info → monai_weekly-1.6.dev2524.dist-info}/RECORD

@@ -1,5 +1,5 @@
-monai/__init__.py,sha256=eZQIXyDYcb7TrjIfc4e844rVm7n1JSyFVCrSaNR63h8,4095
-monai/_version.py,sha256=7wXHOrAWUBKXGlWLCp1a8MS05DKSCeFNQSp2QC6OdFY,503
+monai/__init__.py,sha256=Tuhv3OtRt0Z1-fcKWtwVyLxG1O1LkQDCpUwdIa2LHYU,4095
+monai/_version.py,sha256=E7VRkhmkHHNCXgdYbiP8HabZyg5eJjmgwqx2qSkoFDY,503
 monai/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 monai/_extensions/__init__.py,sha256=NEBPreRhQ8H9gVvgrLr_y52_TmqB96u_u4VQmeNT93I,642
 monai/_extensions/loader.py,sha256=7SiKw36q-nOzH8CRbBurFrz7GM40GCu7rc93Tm8XpnI,3643
@@ -197,7 +197,7 @@ monai/handlers/utils.py,sha256=Ib1u-PLrtIkiLqTfREnrCWpN4af1btdNzkyMZuuuYyU,10239
 monai/handlers/validation_handler.py,sha256=NZO21c6zzXbmAgJZHkkdoZQSQIHwuxh94QD3PLUldGU,3674
 monai/inferers/__init__.py,sha256=K74t_RCeUPdEZvHzIPzVAwZ9DtmouLqhb3qDEmFBWs4,1107
 monai/inferers/inferer.py,sha256=kfrxbSRAug_2oO943gQNHPfhfllu-DrseimT3vEb7uE,97280
-monai/inferers/merger.py,sha256=JxSLdlXTKW1xug11UWQNi6dNtpqVRbGCLc-ifj06g8U,16613
+monai/inferers/merger.py,sha256=fjoaFq5HuKrZ1A0UwlMtd38fk2coa8axJlOjPfP8DQg,21575
 monai/inferers/splitter.py,sha256=_hTnFdvDNRckkA7ZGQehVsNZw83oXoGFWyk5VXNqgJg,21149
 monai/inferers/utils.py,sha256=dvZBCAjaPa8xXcJuXRzNQ-fBzteauzkKbxE5YZdGBGY,20374
 monai/losses/__init__.py,sha256=igy7BjoQzM3McmJPD2tmeiW2ljSXfB2HBdc4YiDzYEg,1778
@@ -426,7 +426,7 @@ monai/visualize/img2tensorboard.py,sha256=n4ztSa5BQAUxSTGvi2tp45v-F7-RNgSlbsdy-9
 monai/visualize/occlusion_sensitivity.py,sha256=0SwhLO7ePDfIXJj67_UmXDZLxXItMeM-uNrPaCE0xXg,18159
 monai/visualize/utils.py,sha256=B-MhTVs7sQbIqYS3yPnpBwPw2K82rE2PBtGIfpwZtWM,9894
 monai/visualize/visualizer.py,sha256=qckyaMZCbezYUwE20k5yc-Pb7UozVavMDbrmyQwfYHY,1377
-monai_weekly-1.5.dev2522.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+monai_weekly-1.6.dev2524.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
 tests/apps/__init__.py,sha256=s9djSd6kvViPnFvMR11Dgd30Lv4oY6FaPJr4ZZJZLq0,573
 tests/apps/test_auto3dseg_bundlegen.py,sha256=FpTJo9Lfe8vdhGuWeZ9y1BQmqYwTt-s8mDVtoLGAz_I,5594
 tests/apps/test_check_hash.py,sha256=MuZslW2DDCxHKEo6-PiL7hnbxGuZRRYf6HOh3ZQv1qQ,1761
@@ -584,7 +584,7 @@ tests/inferers/test_slice_inferer.py,sha256=kzaJjjTnf2rAiR75l8A_J-Kie4NaLj2bogi-
 tests/inferers/test_sliding_window_inference.py,sha256=6DiCbVhsGVbicepb6rJ-9ThFY3YToRb4E5qh8mJew5I,15529
 tests/inferers/test_sliding_window_splitter.py,sha256=LXCqtyMx8D2Ps29EeoiEZTBsBsTYN6H4pBTMwT8Ru0I,9620
 tests/inferers/test_wsi_sliding_window_splitter.py,sha256=Y0d1Er4kOljfZ79pdd2U2I59M6l3xYuvfvEL0T_wIw8,8434
-tests/inferers/test_zarr_avg_merger.py,sha256=CRC_lzwzafOfV3RP0DvuNorQiEnIV0fjMqvHjgFhIbQ,10594
+tests/inferers/test_zarr_avg_merger.py,sha256=o4civt2BgOzI2P2KFD5LVLlw3AVsygXVbgild-fd4pM,15421
 tests/integration/__init__.py,sha256=s9djSd6kvViPnFvMR11Dgd30Lv4oY6FaPJr4ZZJZLq0,573
 tests/integration/test_auto3dseg_ensemble.py,sha256=CIdgRSNX1VBfSmSvx_8HgsIfFsJpoLAC5sg8_HrRP_c,7570
 tests/integration/test_auto3dseg_hpo.py,sha256=J5Us-7iVE5UFbbsZdzNVcvPPQENTlPFarVsSt7KR3Yo,7286
@@ -592,7 +592,7 @@ tests/integration/test_deepedit_interaction.py,sha256=tmryp1cP_QlI_tgguZybRZc7-F
 tests/integration/test_hovernet_nuclear_type_post_processingd.py,sha256=yTRmYdQBXEMMmXJjPDBPMxPSkLWj2U3bdRhaAfDXrpE,2661
 tests/integration/test_integration_autorunner.py,sha256=tDK1XkMZp4hehfuzMr2LQIgavP36L_vkFcOcI1Z68Lk,7571
 tests/integration/test_integration_bundle_run.py,sha256=uO87WnnG3EYnAxhudpfHy7fyxHNNzifFTw2rrMm_6XU,10734
-tests/integration/test_integration_classification_2d.py,sha256=psUvLWNtndkPkgc14YCKqvVQJ9oS1EBdxpg3dOqoF7E,11373
+tests/integration/test_integration_classification_2d.py,sha256=4NT_dkty3cpSvzPGcz9l4_e7uhAsWOwNEsXqz0XEc5c,11373
 tests/integration/test_integration_determinism.py,sha256=AiSBXHcPwDtKRbt_lejI-IDDkYtDWccMkNVoHuyrtU0,3172
 tests/integration/test_integration_fast_train.py,sha256=WxEIJV52F0Cf2wmGlIQDiVs1m2QZrvxmta_UAsa0OCI,9736
 tests/integration/test_integration_gpu_customization.py,sha256=z-w6iBaY72LEi8TBVxZuzvsEBgBecZAP2YPwl6KFUhA,5547
@@ -605,7 +605,7 @@ tests/integration/test_integration_stn.py,sha256=1bwzCn8X-1xjV-SGalOtlpRPLFnYpDG
 tests/integration/test_integration_unet_2d.py,sha256=rMOCG7eYt3jrCjG5HXCfTdF-XnCvyO1X631H7l-F1w4,2376
 tests/integration/test_integration_workers.py,sha256=bSHRXskT-FSRXg5GGj30mIYRPgG_wtVWfcc11XwpsEE,2195
 tests/integration/test_integration_workflows.py,sha256=B3INF2PzXwggxxEHGCaXGH754htlsvJbU7Dw1J73Sg0,13971
-tests/integration/test_integration_workflows_adversarial.py,sha256=q0QvEOGf97z0MKyoqHY3iwTP3jsk9NHiSwGtiAaNhPo,6375
+tests/integration/test_integration_workflows_adversarial.py,sha256=6PAjwQHePPmL5eW7RpSRVtYFvWD5n-vOLhi3a9S7sJc,6375
 tests/integration/test_integration_workflows_gan.py,sha256=42cSwfIgXHmXqBF5a5tnWJV5NT4nfJzn2WSw6GguCRY,5496
 tests/integration/test_loader_semaphore.py,sha256=-Fk4qRA6Up1Lf2uK79Y8WRmDIlOPJEU01kb-3cRX0aM,1195
 tests/integration/test_mapping_filed.py,sha256=19C_tQ5YvYI0b8v-QsHF1sWDymq5D0chKoe86JICREo,4813
@@ -1189,7 +1189,7 @@ tests/visualize/test_vis_gradcam.py,sha256=WpA-pvTB75eZs7JoIc5qyvOV9PwgkzWI8-Vow
 tests/visualize/utils/__init__.py,sha256=s9djSd6kvViPnFvMR11Dgd30Lv4oY6FaPJr4ZZJZLq0,573
 tests/visualize/utils/test_blend_images.py,sha256=RVs2p_8RWQDfhLHDNNtZaMig27v8o0km7XxNa-zWjKE,2274
 tests/visualize/utils/test_matshow3d.py,sha256=wXYj77L5Jvnp0f6DvL1rsi_-YlCxS0HJ9hiPmrbpuP8,5021
-monai_weekly-1.5.dev2522.dist-info/METADATA,sha256=SECyXqQSdsULPcpimKHFcaPJko0nMWSVYjg31Za4OHU,13152
-monai_weekly-1.5.dev2522.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-monai_weekly-1.5.dev2522.dist-info/top_level.txt,sha256=hn2Y6P9xBf2R8faMeVMHhPMvrdDKxMsIOwMDYI0yTjs,12
-monai_weekly-1.5.dev2522.dist-info/RECORD,,
+monai_weekly-1.6.dev2524.dist-info/METADATA,sha256=B9GjloAyVPFRuOtOEfOkjrECaE6O6Ij2c-a4sja7jmA,13285
+monai_weekly-1.6.dev2524.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+monai_weekly-1.6.dev2524.dist-info/top_level.txt,sha256=hn2Y6P9xBf2R8faMeVMHhPMvrdDKxMsIOwMDYI0yTjs,12
+monai_weekly-1.6.dev2524.dist-info/RECORD,,

tests/inferers/test_zarr_avg_merger.py

@@ -24,6 +24,7 @@ from tests.test_utils import assert_allclose
 
 np.seterr(divide="ignore", invalid="ignore")
 zarr, has_zarr = optional_import("zarr")
+print(version_geq(get_package_version("zarr"), "3.0.0"))
 if has_zarr:
     if version_geq(get_package_version("zarr"), "3.0.0"):
         directory_store = zarr.storage.LocalStore("test.zarr")
@@ -200,9 +201,20 @@ TEST_CASE_12_CHUNKS = [
     TENSOR_4x4,
 ]
 
-# test for LZ4 compressor
+# Define zarr v3 codec configurations with proper bytes codec
+ZARR_V3_LZ4_CODECS = [{"name": "bytes", "configuration": {}}, {"name": "blosc", "configuration": {"cname": "lz4"}}]
+
+ZARR_V3_PICKLE_CODECS = [{"name": "bytes", "configuration": {}}, {"name": "blosc", "configuration": {"cname": "zstd"}}]
+
+ZARR_V3_LZMA_CODECS = [{"name": "bytes", "configuration": {}}, {"name": "blosc", "configuration": {"cname": "zlib"}}]
+
+# test for LZ4 compressor (zarr v2) or codecs (zarr v3)
 TEST_CASE_13_COMPRESSOR_LZ4 = [
-    dict(merged_shape=TENSOR_4x4.shape, compressor="LZ4"),
+    (
+        dict(merged_shape=TENSOR_4x4.shape, compressor="LZ4")
+        if not version_geq(get_package_version("zarr"), "3.0.0")
+        else dict(merged_shape=TENSOR_4x4.shape, codecs=ZARR_V3_LZ4_CODECS)
+    ),
     [
         (TENSOR_4x4[..., :2, :2], (0, 0)),
         (TENSOR_4x4[..., :2, 2:], (0, 2)),
@@ -212,9 +224,13 @@ TEST_CASE_13_COMPRESSOR_LZ4 = [
     TENSOR_4x4,
 ]
 
-# test for pickle compressor
+# test for pickle compressor (zarr v2) or codecs (zarr v3)
 TEST_CASE_14_COMPRESSOR_PICKLE = [
-    dict(merged_shape=TENSOR_4x4.shape, compressor="Pickle"),
+    (
+        dict(merged_shape=TENSOR_4x4.shape, compressor="Pickle")
+        if not version_geq(get_package_version("zarr"), "3.0.0")
+        else dict(merged_shape=TENSOR_4x4.shape, codecs=ZARR_V3_PICKLE_CODECS)
+    ),
     [
         (TENSOR_4x4[..., :2, :2], (0, 0)),
         (TENSOR_4x4[..., :2, 2:], (0, 2)),
@@ -224,9 +240,13 @@ TEST_CASE_14_COMPRESSOR_PICKLE = [
     TENSOR_4x4,
 ]
 
-# test for LZMA compressor
+# test for LZMA compressor (zarr v2) or codecs (zarr v3)
 TEST_CASE_15_COMPRESSOR_LZMA = [
-    dict(merged_shape=TENSOR_4x4.shape, compressor="LZMA"),
+    (
+        dict(merged_shape=TENSOR_4x4.shape, compressor="LZMA")
+        if not version_geq(get_package_version("zarr"), "3.0.0")
+        else dict(merged_shape=TENSOR_4x4.shape, codecs=ZARR_V3_LZMA_CODECS)
+    ),
     [
         (TENSOR_4x4[..., :2, :2], (0, 0)),
         (TENSOR_4x4[..., :2, 2:], (0, 2)),
@@ -260,6 +280,48 @@ TEST_CASE_17_WITHOUT_LOCK = [
     TENSOR_4x4,
 ]
 
+# test with codecs for zarr v3
+TEST_CASE_18_CODECS = [
+    dict(merged_shape=TENSOR_4x4.shape, codecs=ZARR_V3_LZ4_CODECS),
+    [
+        (TENSOR_4x4[..., :2, :2], (0, 0)),
+        (TENSOR_4x4[..., :2, 2:], (0, 2)),
+        (TENSOR_4x4[..., 2:, :2], (2, 0)),
+        (TENSOR_4x4[..., 2:, 2:], (2, 2)),
+    ],
+    TENSOR_4x4,
+]
+
+# test with value_codecs for zarr v3
+TEST_CASE_19_VALUE_CODECS = [
+    dict(
+        merged_shape=TENSOR_4x4.shape,
+        value_codecs=[{"name": "bytes", "configuration": {}}, {"name": "blosc", "configuration": {"cname": "zstd"}}],
+    ),
+    [
+        (TENSOR_4x4[..., :2, :2], (0, 0)),
+        (TENSOR_4x4[..., :2, 2:], (0, 2)),
+        (TENSOR_4x4[..., 2:, :2], (2, 0)),
+        (TENSOR_4x4[..., 2:, 2:], (2, 2)),
+    ],
+    TENSOR_4x4,
+]
+
+# test with count_codecs for zarr v3
+TEST_CASE_20_COUNT_CODECS = [
+    dict(
+        merged_shape=TENSOR_4x4.shape,
+        count_codecs=[{"name": "bytes", "configuration": {}}, {"name": "blosc", "configuration": {"cname": "zlib"}}],
+    ),
+    [
+        (TENSOR_4x4[..., :2, :2], (0, 0)),
+        (TENSOR_4x4[..., :2, 2:], (0, 2)),
+        (TENSOR_4x4[..., 2:, :2], (2, 0)),
+        (TENSOR_4x4[..., 2:, 2:], (2, 2)),
+    ],
+    TENSOR_4x4,
+]
+
 ALL_TESTS = [
     TEST_CASE_0_DEFAULT_DTYPE,
     TEST_CASE_1_DEFAULT_DTYPE,
@@ -276,11 +338,15 @@ ALL_TESTS = [
     TEST_CASE_12_CHUNKS,
     TEST_CASE_16_WITH_LOCK,
     TEST_CASE_17_WITHOUT_LOCK,
+    # Add compression/codec tests regardless of zarr version - they're now version-aware
+    TEST_CASE_13_COMPRESSOR_LZ4,
+    TEST_CASE_14_COMPRESSOR_PICKLE,
+    TEST_CASE_15_COMPRESSOR_LZMA,
 ]
 
-# add compression tests only when using Zarr version before 3.0
-if not version_geq(get_package_version("zarr"), "3.0.0"):
-    ALL_TESTS += [TEST_CASE_13_COMPRESSOR_LZ4, TEST_CASE_14_COMPRESSOR_PICKLE, TEST_CASE_15_COMPRESSOR_LZMA]
+# Add zarr v3 specific codec tests only when using Zarr version 3.0 or later
+if version_geq(get_package_version("zarr"), "3.0.0"):
+    ALL_TESTS += [TEST_CASE_18_CODECS, TEST_CASE_19_VALUE_CODECS, TEST_CASE_20_COUNT_CODECS]
 
 
 @unittest.skipUnless(has_zarr and has_numcodecs, "Requires zarr (and numcodecs) packages.)")
@@ -288,16 +354,57 @@ class ZarrAvgMergerTests(unittest.TestCase):
 
     @parameterized.expand(ALL_TESTS)
     def test_zarr_avg_merger_patches(self, arguments, patch_locations, expected):
+        is_zarr_v3 = version_geq(get_package_version("zarr"), "3.0.0")
         codec_reg = numcodecs.registry.codec_registry
-        if "compressor" in arguments:
-            if arguments["compressor"] != "default":
+
+        # Handle compressor/codecs based on zarr version
+        if "compressor" in arguments and is_zarr_v3:
+            # For zarr v3, convert compressor to codecs
+            if arguments["compressor"] != "default" and arguments["compressor"] is not None:
+                compressor_name = arguments["compressor"].lower()
+                if compressor_name == "lz4":
+                    arguments["codecs"] = ZARR_V3_LZ4_CODECS
+                elif compressor_name == "pickle":
+                    arguments["codecs"] = ZARR_V3_PICKLE_CODECS
+                elif compressor_name == "lzma":
+                    arguments["codecs"] = ZARR_V3_LZMA_CODECS
+                # Remove compressor as it's not supported in zarr v3
+                del arguments["compressor"]
+        elif "compressor" in arguments and not is_zarr_v3:
+            # For zarr v2, use the compressor registry
+            if arguments["compressor"] != "default" and arguments["compressor"] is not None:
                 arguments["compressor"] = codec_reg[arguments["compressor"].lower()]()
-        if "value_compressor" in arguments:
-            if arguments["value_compressor"] != "default":
+
+        # Same for value_compressor
+        if "value_compressor" in arguments and is_zarr_v3:
+            if arguments["value_compressor"] != "default" and arguments["value_compressor"] is not None:
+                compressor_name = arguments["value_compressor"].lower()
+                if compressor_name == "lz4":
+                    arguments["value_codecs"] = ZARR_V3_LZ4_CODECS
+                elif compressor_name == "pickle":
+                    arguments["value_codecs"] = ZARR_V3_PICKLE_CODECS
+                elif compressor_name == "lzma":
+                    arguments["value_codecs"] = ZARR_V3_LZMA_CODECS
+                del arguments["value_compressor"]
+        elif "value_compressor" in arguments and not is_zarr_v3:
+            if arguments["value_compressor"] != "default" and arguments["value_compressor"] is not None:
                 arguments["value_compressor"] = codec_reg[arguments["value_compressor"].lower()]()
-        if "count_compressor" in arguments:
-            if arguments["count_compressor"] != "default":
+
+        # Same for count_compressor
+        if "count_compressor" in arguments and is_zarr_v3:
+            if arguments["count_compressor"] != "default" and arguments["count_compressor"] is not None:
+                compressor_name = arguments["count_compressor"].lower()
+                if compressor_name == "lz4":
+                    arguments["count_codecs"] = ZARR_V3_LZ4_CODECS
+                elif compressor_name == "pickle":
+                    arguments["count_codecs"] = ZARR_V3_PICKLE_CODECS
+                elif compressor_name == "lzma":
+                    arguments["count_codecs"] = ZARR_V3_LZMA_CODECS
+                del arguments["count_compressor"]
+        elif "count_compressor" in arguments and not is_zarr_v3:
+            if arguments["count_compressor"] != "default" and arguments["count_compressor"] is not None:
                 arguments["count_compressor"] = codec_reg[arguments["count_compressor"].lower()]()
+
         merger = ZarrAvgMerger(**arguments)
         for pl in patch_locations:
             merger.aggregate(pl[0], pl[1])
@@ -320,7 +427,3 @@ class ZarrAvgMergerTests(unittest.TestCase):
     def test_zarr_avg_merge_none_merged_shape_error(self):
         with self.assertRaises(ValueError):
             ZarrAvgMerger(merged_shape=None)
-
-
-if __name__ == "__main__":
-    unittest.main()

tests/integration/test_integration_classification_2d.py

@@ -256,7 +256,7 @@ class IntegrationClassification2D(DistTestCase):
         # check training properties
         self.assertTrue(test_integration_value(TASK, key="losses", data=losses, rtol=1e-2))
         self.assertTrue(test_integration_value(TASK, key="best_metric", data=best_metric, rtol=1e-4))
-        np.testing.assert_allclose(best_metric_epoch, 4)
+        np.testing.assert_allclose(best_metric_epoch, 1)
         model_file = os.path.join(self.data_dir, "best_metric_model.pth")
         self.assertTrue(os.path.exists(model_file))
         # check inference properties
@@ -268,7 +268,7 @@ class IntegrationClassification2D(DistTestCase):
 
     def test_training(self):
         repeated = []
-        for i in range(1):
+        for i in range(2):
             results = self.train_and_infer(i)
             repeated.append(results)
         np.testing.assert_allclose(repeated[0], repeated[1])

tests/integration/test_integration_workflows_adversarial.py

@@ -158,7 +158,7 @@ class IntegrationWorkflowsAdversarialTrainer(DistTestCase):
         set_determinism(seed=None)
         shutil.rmtree(self.data_dir)
 
-    @TimedCall(seconds=200, daemon=False)
+    @TimedCall(seconds=300, daemon=False)
     def test_training(self):
         torch.manual_seed(0)