def _batch_zip_sub_data_containers(self, data_container: DataContainer):
"""
Zip sub data containers on the batch dimension.
:param data_container: data container to zip
:type data_container: DataContainer
:return: base step, data container
:rtype: DataContainer
"""
sub_data_containers_to_zip = []
if self.data_sources is None:
self.data_sources = data_container.get_sub_data_container_names()
for name, sub_data_container in data_container.sub_data_containers:
if name in self.data_sources:
sub_data_containers_to_zip.append(sub_data_container)
for data_container_to_zip in sub_data_containers_to_zip:
data_container = self._batch_zip_sub_data_container(data_container, data_container_to_zip)
return data_container
def _concatenate_sub_data_containers(self, data_container: DataContainer):
"""
Merge sub data containers into the current data container.
:param data_container: data container to zip
:type data_container: DataContainer
:return: base step, data container
:rtype: DataContainer
"""
sub_data_containers_to_zip = []
if self.data_sources is None:
self.data_sources = data_container.get_sub_data_container_names()
for name, sub_data_container in data_container.sub_data_containers:
if name in self.data_sources:
sub_data_containers_to_zip.append(sub_data_container)
for data_container_to_zip in sub_data_containers_to_zip:
data_container = self._concatenate_sub_data_container(data_container, data_container_to_zip)
return data_container
