mlx - trainer update for compilation and data adapters (#20957)

* mlx trainer update and data adapters

* formatting

Author: acsweet

keras/src/backend/common/dtypes_test.py

@@ -47,10 +47,15 @@ def setUp(self):
 
         self.jax_enable_x64 = enable_x64()
         self.jax_enable_x64.__enter__()
+        if backend.backend() == "mlx":
+            self.mlx_cpu_context = backend.core.enable_float64()
+            self.mlx_cpu_context.__enter__()
         return super().setUp()
 
     def tearDown(self):
         self.jax_enable_x64.__exit__(None, None, None)
+        if backend.backend() == "mlx":
+            self.mlx_cpu_context.__exit__(None, None, None)
         return super().tearDown()
 
     @parameterized.named_parameters(

keras/src/backend/common/variables_test.py

@@ -256,6 +256,13 @@ def test_standardize_dtype(self, dtype):
         ):
             self.skipTest(f"openvino backend does not support dtype {dtype}")
 
+        if backend.backend() == "mlx" and dtype in (
+            "complex128",
+            "float8_e4m3fn",
+            "float8_e5m2",
+        ):
+            self.skipTest(f"mlx backend does not support dtype {dtype}")
+
         x = backend.convert_to_tensor(np.zeros(()), dtype)
         actual = standardize_dtype(x.dtype)
         self.assertEqual(actual, dtype)
@@ -805,6 +812,8 @@ def test_invalid_float(self):
 elif backend.backend() == "openvino":
     # TODO: openvino doesn't support complex
     ALL_DTYPES = [x for x in ALL_DTYPES if x not in ["complex128", "complex64"]]
+elif backend.backend() == "mlx":
+    ALL_DTYPES = [x for x in ALL_DTYPES if x not in ["complex128"]]
 # Remove float8 dtypes for the following tests
 ALL_DTYPES = [x for x in ALL_DTYPES if x not in dtypes.FLOAT8_TYPES]
 NON_COMPLEX_DTYPES = [x for x in ALL_DTYPES if x and x not in COMPLEX_DTYPES]
@@ -818,10 +827,15 @@ def setUp(self):
 
         self.jax_enable_x64 = enable_x64()
         self.jax_enable_x64.__enter__()
+        if backend.backend() == "mlx":
+            self.mlx_cpu_context = backend.core.enable_float64()
+            self.mlx_cpu_context.__enter__()
         return super().setUp()
 
     def tearDown(self):
         self.jax_enable_x64.__exit__(None, None, None)
+        if backend.backend() == "mlx":
+            self.mlx_cpu_context.__exit__(None, None, None)
         return super().tearDown()
 
     @parameterized.named_parameters(

keras/src/backend/mlx/__init__.py

@@ -17,6 +17,7 @@
 from keras.src.backend.mlx.core import cond
 from keras.src.backend.mlx.core import convert_to_numpy
 from keras.src.backend.mlx.core import convert_to_tensor
+from keras.src.backend.mlx.core import device_scope
 from keras.src.backend.mlx.core import is_tensor
 from keras.src.backend.mlx.core import random_seed_dtype
 from keras.src.backend.mlx.core import scatter

keras/src/backend/mlx/core.py

@@ -25,7 +25,7 @@
 MLX_DTYPES = {
     "float16": mx.float16,
     "float32": mx.float32,
-    "float64": None,  # mlx only supports float64 on cpu
+    "float64": mx.float64,  # for mlx float64 only supported on cpu
     "uint8": mx.uint8,
     "uint16": mx.uint16,
     "uint32": mx.uint32,
@@ -104,19 +104,7 @@ def convert_to_tensor(x, dtype=None, sparse=None, ragged=None):
         return mx.array(x, dtype=mlx_dtype)
 
     if isinstance(x, list):
-
-        def to_scalar_list(x):
-            if isinstance(x, list):
-                return [to_scalar_list(xi) for xi in x]
-            elif isinstance(x, mx.array):
-                if x.ndim == 0:
-                    return x.item()
-                else:
-                    return x.tolist()
-            else:
-                return x
-
-        return mx.array(to_scalar_list(x), dtype=mlx_dtype)
+        return mx.array(x, dtype=mlx_dtype)
 
     if _is_h5py_dataset(x):
         if h5py is None:
@@ -592,3 +580,41 @@ def __init__(self, fun):
     def __call__(self, *args, **kwargs):
         outputs, _ = self.fun(*args, **kwargs)
         return outputs
+
+
+def enable_float64():
+    """Returns context manager forcing operations on cpu
+
+    MLX requires operations involving `float64` to be on cpu,
+    mimicking jax's `enable_x64()`
+
+    Usage:
+    ```
+        a = mx.array([1, 2, 3], dtype=mx.float64)
+        b = mx.array([4, 5, 6], dtype=mx.float64)
+
+        with enable_float64():
+            c = mx.add(a, b)
+
+        # OR
+        mlx_cpu_context = mx.stream(mx.cpu)
+        mlx_cpu_context.__enter__()
+        c = mx.add(a, b)
+        mlx_cpu_context.__exit__(None, None, None)
+    ```
+    """
+    return mx.stream(mx.cpu)
+
+
+def device_scope(device_name):
+    if isinstance(device_name, str):
+        mlx_device = mx.cpu if "cpu" in device_name.lower() else mx.gpu
+    elif not isinstance(device_name, mx.Device):
+        raise ValueError(
+            "Invalid value for argument `device_name`. "
+            "Expected a string like 'gpu:0' or a `mlx.core.Device` instance. "
+            f"Received: device_name='{device_name}'"
+        )
+    else:
+        mlx_device = device_name
+    return mx.stream(mlx_device)

keras/src/backend/mlx/numpy.py

@@ -291,10 +291,16 @@ def right_shift(x, y):
 
 
 def _bincount_1d(x, weights=None, minlength=0):
-    length = builtins.max(builtins.max(x) + 1, minlength or 0)
+    x_max = mx.max(x)
+    length = mx.maximum(x_max + 1, minlength)
+
     counts = mx.zeros(length)
-    w = weights if weights is not None else mx.ones_like(x)
-    return counts.at[x].add(w)
+    if weights is None:
+        counts = counts.at[x].add(1)
+    else:
+        counts = counts.at[x].add(weights)
+
+    return counts
 
 
 def bincount(x, weights=None, minlength=0, sparse=False):

keras/src/backend/mlx/random.py

@@ -226,7 +226,7 @@ def _sample_gamma(shape, a, b, key):
 def binomial(shape, counts, probabilities, dtype=None, seed=None):
     # Binomial(n, p) distribution by summing n Bernoulli(p) samples
     dtype = to_mlx_dtype(dtype)
-    key = mx.random.key(seed)
+    key = mlx_draw_seed(seed)
 
     if isinstance(shape, int):
         shape = (shape,)