returnn 1.20250416.132454__tar.gz → 1.20250416.172956__tar.gz

{returnn-1.20250416.132454/returnn.egg-info → returnn-1.20250416.172956}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20250416.132454
+Version: 1.20250416.172956
 Summary: The RWTH extensible training framework for universal recurrent neural networks
 Home-page: https://github.com/rwth-i6/returnn/
 Author: Albert Zeyer

returnn-1.20250416.172956/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20250416.172956'
+long_version = '1.20250416.172956+git.505963b'

{returnn-1.20250416.132454 → returnn-1.20250416.172956}/returnn/datasets/basic.py

@@ -864,18 +864,16 @@ class Dataset:
             data = self.get_data(seq_idx, key)
             return data[s0_start:s0_end]
 
-    def get_tag(self, sorted_seq_idx):
+    def get_tag(self, sorted_seq_idx: int) -> str:
         """
-        :param int sorted_seq_idx:
-        :rtype: str
+        :param sorted_seq_idx:
         """
         return "seq-%i" % sorted_seq_idx
 
-    def get_all_tags(self):
+    def get_all_tags(self) -> List[str]:
         """
         :return: list of all seq tags, of the whole dataset, without partition epoch.
           Note that this is not possible with all datasets.
-        :rtype: list[str]
         """
         raise OptionalNotImplementedError(f"{self} get_all_tags not implemented")
 

{returnn-1.20250416.132454 → returnn-1.20250416.172956}/returnn/datasets/hdf.py

@@ -1201,6 +1201,7 @@ class SimpleHDFWriter:
             shape = [None] * ndim  # type: typing.List[typing.Optional[int]]
             if ndim >= 2:
                 shape[-1] = dim
+            assert all(shape[1:]), f"{self} extra {data_key!r} supports only dyn dim in first axis, got shape {shape!r}"
             if dtype == "string":
                 # noinspection PyUnresolvedReferences
                 dtype = h5py.special_dtype(vlen=str)
@@ -1237,10 +1238,15 @@ class SimpleHDFWriter:
             self._datasets[name] = self._file.create_dataset(
                 name, raw_data.shape, raw_data.dtype, maxshape=tuple(None for _ in raw_data.shape)
             )
+            expected_shape = (raw_data.shape[0],) + self._datasets[name].shape[1:]
         else:
             old_shape = self._datasets[name].shape
             self._datasets[name].resize(old_shape[0] + raw_data.shape[0], axis=0)
+            expected_shape = (raw_data.shape[0],) + old_shape[1:]
         # append raw data to dataset
+        assert (
+            expected_shape == raw_data.shape
+        ), f"{self} insert: shape mismatch: expected {expected_shape}, got {raw_data.shape}"
         self._datasets[name][self._file.attrs["numTimesteps"] :] = raw_data
         self._file.attrs["numTimesteps"] += raw_data.shape[0]
         self._file.attrs["numSeqs"] += 1
@@ -1286,13 +1292,17 @@ class SimpleHDFWriter:
                 self._seq_lengths[seq_idx, data_key_idx_0 + 1] = self._extra_num_time_steps[data_key_]
 
         self._extra_num_time_steps[data_key] += raw_data.shape[0]
-        self._datasets[data_key].resize(self._extra_num_time_steps[data_key], axis=0)
+        hdf_data = self._datasets[data_key]
+        hdf_data.resize(self._extra_num_time_steps[data_key], axis=0)
 
         data_key_idx = sorted(self._prepared_extra).index(data_key) + 1
         self._seq_lengths[seq_idx, data_key_idx] = raw_data.shape[0]
 
         offset = self._extra_num_time_steps[data_key] - raw_data.shape[0]
-        hdf_data = self._datasets[data_key]
+        expected_shape = (raw_data.shape[0],) + hdf_data.shape[1:]
+        assert (
+            expected_shape == raw_data.shape
+        ), f"{self} insert other {data_key!r}: shape mismatch: expected {expected_shape}, got {raw_data.shape}"
         hdf_data[offset:] = raw_data
 
     def insert_batch(self, inputs, seq_len, seq_tag, extra=None):

{returnn-1.20250416.132454 → returnn-1.20250416.172956}/returnn/datasets/postprocessing.py

@@ -243,6 +243,13 @@ class PostprocessingDataset(CachedDataset2):
         assert self._dataset is not None
         return self._dataset.get_total_num_seqs(fast=fast)
 
+    def get_all_tags(self) -> List[str]:
+        """:return: all tags"""
+        if self._map_seq_stream is not None:
+            raise util.OptionalNotImplementedError(f"{self}: get_all_tags not allowed when map_seq_stream is set.")
+        assert self._dataset is not None
+        return self._dataset.get_all_tags()
+
     def supports_sharding(self) -> bool:
         """:return: whether this dataset supports sharding"""
         assert self._dataset is not None

{returnn-1.20250416.132454 → returnn-1.20250416.172956/returnn.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20250416.132454
+Version: 1.20250416.172956
 Summary: The RWTH extensible training framework for universal recurrent neural networks
 Home-page: https://github.com/rwth-i6/returnn/
 Author: Albert Zeyer

{returnn-1.20250416.132454 → returnn-1.20250416.172956}/tests/test_HDFDataset.py

@@ -576,6 +576,81 @@ def test_SimpleHDFWriter_labels():
         assert reader.seq_lens[i]["data"] == seq_len
 
 
+def test_SimpleHDFWriter_extra_with_feat():
+    from returnn.tensor import TensorDict, Dim, batch_dim
+    from returnn.tensor.utils import tensor_dict_fill_random_numpy_
+
+    fn = get_test_tmp_file(suffix=".hdf")
+    os.remove(fn)  # SimpleHDFWriter expects that the file does not exist
+
+    # See i6_experiments.users.zeyer.forward_to_hdf._returnn_get_forward_callback.
+    # (Maybe also somewhat similar: HDFDumpLayer)
+    spatial_dim = Dim(None, name="spatial")
+    k_dim = Dim(5, name="k")
+    vocab_dim = Dim(11, name="vocab")
+    expected_outputs = TensorDict()
+    expected_outputs.update(
+        {
+            "output": {"dims": [batch_dim, spatial_dim, k_dim], "dtype": "int32", "sparse_dim": vocab_dim},
+            "k_log_probs": {"dims": [batch_dim, spatial_dim, k_dim], "dtype": "float32"},
+        },
+        auto_convert=True,
+    )
+    output = expected_outputs["output"]
+
+    writer = SimpleHDFWriter(
+        filename=fn,
+        dim=output.dim,
+        ndim=output.ndim,
+        labels=output.vocab and output.vocab.labels,
+        # Note: in HDFDumpLayer, we use ndim=min(v.ndim - len(v.size_placeholder) + 1, 2) instead of just v.ndim...
+        extra_type={k: (v.shape[-1], v.ndim, v.dtype) for k, v in expected_outputs.data.items() if k != "output"},
+        extra_labels={k: v.vocab.labels for k, v in expected_outputs.data.items() if k != "output" and v.vocab},
+    )
+
+    expected_outputs.reset_content()
+    tensor_dict_fill_random_numpy_(expected_outputs)
+
+    # Note: HDFDumpLayer does a lot of extra stuff for the extra data...
+    # We are following i6_experiments.users.zeyer.forward_to_hdf here.
+    # We iterate over the batch dim, to insert each seq separately,
+    # which is important that the extra data seq lens are correct
+    # (insert_batch would not know this otherwise).
+    batch_size = output.raw_tensor.shape[0]
+    for b in range(batch_size):
+        writer.insert_batch(
+            inputs=output.raw_tensor[b : b + 1, : spatial_dim.dyn_size[b]],
+            # We are following i6_experiments.users.zeyer.forward_to_hdf here.
+            # This logic then in turn will flatten the dims internally.
+            seq_len={0: spatial_dim.dyn_size[b : b + 1], 1: [k_dim.dimension]},
+            seq_tag=[f"seq-{b}"],
+            extra={
+                k: v.raw_tensor[b : b + 1, : spatial_dim.dyn_size[b]]
+                for k, v in expected_outputs.data.items()
+                if k != "output"
+            },
+        )
+
+    writer.close()
+
+    dataset = HDFDataset(files=[fn])
+    reader = DatasetTestReader(dataset=dataset)
+    reader.read_all()
+    assert reader.num_seqs > 0
+    assert set(reader.data.keys()) == {"data", "k_log_probs", "sizes"}
+    for i in range(reader.num_seqs):
+        # data dims will be flattened.
+        data0 = reader.data["data"][i]
+        print(f"** seq_idx {i} data key 'data' shape {data0.shape}")
+        assert data0.ndim == 1 and data0.shape[0] % k_dim.dimension == 0
+        time0 = data0.shape[0] // k_dim.dimension
+        sizes0 = reader.data["sizes"][i]
+        assert sizes0.shape == (2,)
+        assert sizes0.tolist() == [time0, k_dim.dimension]
+        probs0 = reader.data["k_log_probs"][i]
+        assert probs0.shape == (time0, k_dim.dimension)
+
+
 class Old2018HDFDataset(CachedDataset):
     """
     Copied and adapted from an early RETURNN version:

returnn-1.20250416.132454/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20250416.132454'
-long_version = '1.20250416.132454+git.6addb6a'