Generic ndarray transformer

Author: MasterSkepticista

openfl/pipelines/no_compression_pipeline.py

@@ -4,12 +4,12 @@
 
 """NoCompressionPipeline module."""
 
-from openfl.pipelines.pipeline import Float32NumpyArrayToBytes, TransformationPipeline
+from openfl.pipelines.pipeline import NumpyArrayToBytes, TransformationPipeline
 
 
 class NoCompressionPipeline(TransformationPipeline):
     """The data pipeline without any compression."""
 
     def __init__(self, **kwargs):
         """Initialize."""
-        super().__init__(transformers=[Float32NumpyArrayToBytes()], **kwargs)
+        super().__init__(transformers=[NumpyArrayToBytes()], **kwargs)

openfl/pipelines/pipeline.py

@@ -93,6 +93,44 @@ def backward(self, data, metadata, **kwargs):
         return np.reshape(flat_array, newshape=array_shape, order="C")
 
 
+class NumpyArrayToBytes(Transformer):
+    """Transformer for converting generic Numpy arrays to bytes."""
+
+    def __init__(self):
+        self.lossy = False
+
+    def forward(self, data: np.ndarray, **kwargs):
+        """Convert a Numpy array to bytes.
+
+        Args:
+            data: The Numpy array to be converted.
+            **kwargs: Additional keyword arguments for the conversion.
+
+        Returns:
+            data_bytes: The data converted to bytes.
+            metadata: The metadata for the conversion.
+        """
+        array_shape = data.shape
+        metadata = {"int_list": list(array_shape), "dtype": str(data.dtype)}
+        data_bytes = data.tobytes(order="C")
+        return data_bytes, metadata
+
+    def backward(self, data, metadata, **kwargs):
+        """Convert bytes back to a Numpy array.
+
+        Args:
+            data: The data in bytes.
+            metadata: The metadata for the conversion.
+
+        Returns:
+            The data converted back to a Numpy array.
+        """
+        array_shape = tuple(metadata["int_list"])
+        dtype = np.dtype(metadata["dtype"])
+        flat_array = np.frombuffer(data, dtype=dtype)
+        return np.reshape(flat_array, newshape=array_shape, order="C")
+
+
 class TransformationPipeline:
     """Data Transformer Pipeline Class.
 

openfl/protocols/base.proto

@@ -22,6 +22,7 @@ message MetadataProto {
   map<int32, float> int_to_float = 1;
   repeated int32 int_list = 2;
   repeated bool bool_list = 3;
+  repeated string dtype = 4;
 }
 
 // handles large size data