Feature/667 zarr 3 support

CHANGELOG.md

@@ -163,6 +163,8 @@
     - renaming `PreConditioner` -> `Preconditioner`
     - fit to `TensorOperator`
     [PR #601](https://github.com/aai-institute/pyDVL/pull/601)
+  - Bumped `zarr` dependency to v3
+    [PR #668](https://github.com/aai-institute/pyDVL/pull/668) 
 
 
 ## 0.9.2 - 🏗  Bug fixes, logging improvement

requirements-extras.txt

@@ -2,5 +2,7 @@ dask>=2023.5.0
 distributed>=2023.5.0
 ray>=0.8
 torch>=2.0.0
-zarr>=2.16.1,<3
+zarr>=2,<3; python_version<'3.11'
+numcodecs<0.16.0; python_version<'3.11'
+zarr>=3,<4; python_version>='3.11'
 pymemcache>=3

requirements-notebooks.txt

@@ -6,5 +6,6 @@ pillow==10.4.0
 torch==2.2.0
 torchvision==0.17.0
 transformers==4.44.2
-zarr==2.16.1
--c requirements-constraints.txt
+zarr>=2,<3; python_version<'3.11'
+numcodecs<0.16.0; python_version<'3.11'
+zarr>=3,<4; python_version>='3.11'

requirements-type-checking.txt

@@ -1,4 +1,4 @@
 mypy==1.15.0
 types-tqdm==4.67.0.20241221
 pandas-stubs
--c requirements-constraints.txt
+numpy<2

setup.py

@@ -27,7 +27,9 @@
             "torch>=2.0.0",
             "dask>=2023.5.0",
             "distributed>=2023.5.0",
-            "zarr>=2.16.1,<3",
+            "zarr>=2,<3; python_version<'3.11'",
+            "numcodecs<0.16.0; python_version<'3.11'",
+            "zarr>=3,<4; python_version>='3.11'",
         ],
         "ray": ["ray>=0.8"],
     },

tests/influence/test_influence_calculator.py

@@ -2,10 +2,10 @@
 import numpy as np
 import pytest
 import torch
-import zarr
 from distributed import Client
 from torch import nn
 from torch.utils.data import DataLoader, TensorDataset
+from zarr.storage import MemoryStore
 
 from pydvl.influence import DaskInfluenceCalculator, InfluenceMode
 from pydvl.influence.base_influence_function_model import (
@@ -306,7 +306,7 @@ def test_sequential_calculator(model_and_data, test_case, mocker):
 
     torch_factors = inf_model.influence_factors(x_test, y_test)
 
-    zarr_factors_store = zarr.MemoryStore()
+    zarr_factors_store = MemoryStore()
     seq_factors_from_zarr = seq_factors_lazy_array.to_zarr(
         zarr_factors_store, TorchNumpyConverter(), return_stored=True
     )
@@ -331,7 +331,7 @@ def test_sequential_calculator(model_and_data, test_case, mocker):
     seq_values_from_factors = seq_values_from_factors_lazy_array.compute(
         aggregator=NestedTorchCatAggregator()
     )
-    zarr_values_from_factors_store = zarr.MemoryStore()
+    zarr_values_from_factors_store = MemoryStore()
     seq_values_from_factors_from_zarr = seq_values_from_factors_lazy_array.to_zarr(
         zarr_values_from_factors_store, TorchNumpyConverter(), return_stored=True
     )
@@ -347,7 +347,7 @@ def test_sequential_calculator(model_and_data, test_case, mocker):
     )
     seq_values = seq_values_lazy_array.compute(aggregator=NestedTorchCatAggregator())
 
-    zarr_values_store = zarr.MemoryStore()
+    zarr_values_store = MemoryStore()
     seq_values_from_zarr = seq_values_lazy_array.to_zarr(
         zarr_values_store, TorchNumpyConverter(), return_stored=True
     )