feat(containers): add flag to allow filtering distributed axis

draco/core/containers.py

@@ -2673,6 +2673,7 @@ def copy_datasets_filter(
     axis: str,
     selection: Union[np.ndarray, list, slice],
     exclude_axes: List[str] = None,
+    allow_distributed: bool = False,
 ):
     """Copy datasets while filtering a given axis.
 
@@ -2689,6 +2690,10 @@ def copy_datasets_filter(
     exclude_axes
         An optional set of axes that if a dataset contains one means it will
         not be copied.
+    allow_distributed, optional
+        Allow the filtered axis to be the distributed axis. This is ONLY
+        valid if filtering is occuring on the local rank only, and mainly
+        exists for compatibility
     """
     exclude_axes_set = set(exclude_axes) if exclude_axes else set()
 
@@ -2717,7 +2722,7 @@ def copy_datasets_filter(
 
         if isinstance(item, memh5.MemDatasetDistributed):
 
-            if item.distributed_axis == axis_ind:
+            if (item.distributed_axis == axis_ind) and not allow_distributed:
                 raise RuntimeError(
                     f"Cannot redistristribute dataset={item.name} along "
                     f"axis={axis_ind} as it is distributed."