[BugFix] Propagate create-dataset kwargs to nested PersistentTensorDicts (#1759)

Author: vmoens

tensordict/persistent.py

@@ -1022,6 +1022,7 @@ def _make_nested(self, key, group):
             backend=self.backend,
             # a persisted batch size was already validated at write time
             validate_batch_size=nested_batch_size is None,
+            **self.kwargs,
         )
         # share the backend instance so per-store state (e.g. the consolidated
         # metadata flag of the zarr backend) is tracked once per store
@@ -2015,6 +2016,7 @@ def _set_metadata(self, orig_metadata_container: PersistentTensorDict):
                 batch_size=td.batch_size,
                 device=td.device,
                 backend=self.backend,
+                **self.kwargs,
             )
             self._nested_tensordicts[key]._backend = self._backend
             self._nested_tensordicts[key].names = td._td_dim_names
@@ -2070,6 +2072,7 @@ def _clone(self, recurse: bool = True, newfile=None) -> PersistentTensorDict:
                 backend=self.backend,
                 device=self.device,
                 batch_size=self.batch_size,
+                **self.kwargs,
             )
             clone._nested_tensordicts = nested_tds
             clone._pin_mem = False

test/memmap/test_h5.py

@@ -140,6 +140,28 @@ def check(x, y):
     tree_map(check, td_dict, td_recon_dict)
 
 
+@pytest.mark.skipif(not _has_h5py, reason="h5py not found.")
+def test_kwargs_passthrough_nested(tmpdir):
+    # create_dataset kwargs must reach the leaves of nested tensordicts too
+    # https://github.com/pytorch/tensordict/issues/1758
+    tmpdir = Path(tmpdir)
+    td = TensorDict(
+        {
+            "a": torch.zeros(64, 3),
+            "b": {"c": torch.zeros(64, 5), "d": {"e": torch.zeros(64, 7)}},
+        },
+        batch_size=[64],
+    )
+    td.to_h5(tmpdir / "file.h5", compression="gzip", compression_opts=9)
+    with h5py.File(tmpdir / "file.h5", "r") as f:
+        for key in ("a", "b/c", "b/d/e"):
+            assert f[key].compression == "gzip", key
+            assert f[key].compression_opts == 9, key
+    td_recon = TensorDict.from_h5(tmpdir / "file.h5")
+    for key in (("a",), ("b", "c"), ("b", "d", "e")):
+        assert (td_recon[key] == td[key]).all(), key
+
+
 if __name__ == "__main__":
     args, unknown = argparse.ArgumentParser().parse_known_args()
     pytest.main([__file__, "--capture", "no", "--exitfirst"] + unknown)

test/memmap/test_zarr.py

@@ -321,6 +321,26 @@ def test_kwargs_passthrough(self, tmp_path):
         root = zarr.open_group(str(tmp_path / "s.zarr"), mode="r")
         assert root["a"].chunks == (16, 64)
 
+    def test_kwargs_passthrough_nested(self, tmp_path):
+        # create-kwargs must reach the leaves of nested tensordicts too
+        # https://github.com/pytorch/tensordict/issues/1758
+        from zarr.codecs import ZstdCodec
+
+        td = TensorDict(
+            {
+                "a": torch.zeros(64, 3),
+                "b": {"c": torch.zeros(64, 3), "d": {"e": torch.zeros(64, 3)}},
+            },
+            batch_size=[64],
+        )
+        td.to_zarr(tmp_path / "s.zarr", chunks=(16,), compressors=ZstdCodec(level=3))
+        root = zarr.open_group(str(tmp_path / "s.zarr"), mode="r")
+        for key in ("a", "b/c", "b/d/e"):
+            assert root[key].chunks == (16, 3), key
+            assert root[key].compressors, key
+        td_back = TensorDict.from_zarr(tmp_path / "s.zarr")
+        assert (td_back == td).all()
+
     def test_chunks_heterogeneous_ranks(self, tmp_path):
         # a single chunks spec constrains the leading dims of every leaf,
         # whatever its rank, and leaves non-tensor payloads untouched