mlx - data adapters and saving (#21023)

* updates to data adapters with mlx arrays and other backends and keras saving

* formatting

Author: acsweet

keras/src/backend/mlx/core.py

@@ -153,6 +153,9 @@ def convert_to_tensors(*xs):
 
 def convert_to_numpy(x):
     # Performs a copy. If we want 0-copy we can pass copy=False
+    if isinstance(x, mx.array) and x.dtype == mx.bfloat16:
+        # mlx currently has an error passing bloat16 array to numpy
+        return np.array(x.astype(mx.float32))
     return np.array(x)
 
 

keras/src/saving/saving_api_test.py

@@ -5,6 +5,7 @@
 from absl import logging
 from absl.testing import parameterized
 
+from keras.src import backend
 from keras.src import layers
 from keras.src.models import Sequential
 from keras.src.saving import saving_api
@@ -122,6 +123,11 @@ def get_model(self, dtype=None):
     )
     def test_basic_load(self, dtype):
         """Test basic model loading."""
+        if backend.backend() == "mlx" and dtype == "float64":
+            self.skipTest(
+                "mlx backend does not yet support float64 in random and uniform"
+            )
+
         model = self.get_model(dtype)
         filepath = os.path.join(self.get_temp_dir(), "test_model.keras")
         saving_api.save_model(model, filepath)
@@ -208,6 +214,13 @@ def get_model(self, dtype=None):
     )
     def test_load_keras_weights(self, source_dtype, dest_dtype):
         """Test loading keras weights."""
+        if backend.backend() == "mlx":
+            if source_dtype == "float64" or dest_dtype == "float64":
+                self.skipTest(
+                    "mlx backend does not yet support float64 in "
+                    "random and uniform"
+                )
+
         src_model = self.get_model(dtype=source_dtype)
         filepath = os.path.join(self.get_temp_dir(), "test_weights.weights.h5")
         src_model.save_weights(filepath)

keras/src/saving/serialization_lib_test.py

@@ -6,6 +6,7 @@
 import pytest
 
 import keras
+from keras.src import backend
 from keras.src import ops
 from keras.src import testing
 from keras.src.saving import serialization_lib
@@ -112,7 +113,11 @@ def test_serialize_ellipsis(self):
         self.assertEqual(..., deserialized)
 
     def test_tensors_and_shapes(self):
-        x = ops.random.normal((2, 2), dtype="float64")
+        if backend.backend() == "mlx":
+            # mlx backend does not yet support float64 in normal and uniform
+            x = ops.random.normal((2, 2), dtype="float32")
+        else:
+            x = ops.random.normal((2, 2), dtype="float64")
         obj = {"x": x}
         _, new_obj, _ = self.roundtrip(obj)
         self.assertAllClose(x, new_obj["x"], atol=1e-5)

keras/src/trainers/data_adapters/array_slicing.py

@@ -146,6 +146,26 @@ def __getitem__(self, indices):
 
         return self.array[mx.array(indices)]
 
+    @classmethod
+    def cast(cls, x, dtype):
+        from keras.src.backend.mlx.core import cast
+
+        return cast(x, dtype)
+
+    @classmethod
+    def convert_to_numpy(cls, x):
+        from keras.src.backend.mlx.core import convert_to_numpy
+
+        return convert_to_numpy(x)
+
+    @classmethod
+    def convert_to_jax_compatible(cls, x):
+        return cls.convert_to_numpy(x)
+
+    @classmethod
+    def convert_to_tf_dataset_compatible(cls, x):
+        return cls.convert_to_numpy(x)
+
 
 class TensorflowSliceable(Sliceable):
     def __getitem__(self, indices):

keras/src/trainers/data_adapters/data_adapter_utils.py

@@ -222,6 +222,8 @@ def convert_to_numpy(x):
                 if is_torch_tensor(x):
                     x = x.cpu()
                 x = np.asarray(x)
+            if is_mlx_array(x):
+                x = np.array(x)
         return x
 
     for batch in iterable:

keras/src/trainers/data_adapters/tf_dataset_adapter.py

@@ -62,9 +62,7 @@ def get_mlx_iterator(self):
         def convert_to_mlx(x):
             if isinstance(x, tf.SparseTensor):
                 x = sparse_to_dense(x)
-            # tensorflow supports the buffer protocol
-            # but requires explicit memoryview with mlx
-            return mlx.core.array(memoryview(x))
+            return mlx.core.array(x)
 
         for batch in self._dataset:
             yield tree.map_structure(convert_to_mlx, batch)