mlx - numpy.searchsorted and numpy.historgram implemented (#20927)

* numpy.searchsorted and numpy.historgram

* clean comments

Author: acsweet

.github/workflows/actions.yml

@@ -78,6 +78,10 @@ jobs:
         if: ${{ matrix.backend == 'jax'}}
         run: |
           python integration_tests/jax_custom_fit_test.py
+      - name: Test MLX-specific integrations
+        if: ${{ matrix.backend == 'mlx'}}
+        run: |
+          python integration_tests/mlx_custom_fit_test.py
       - name: Test TF-specific integrations
         if: ${{ matrix.backend == 'tensorflow'}}
         run: |

integration_tests/mlx_custom_fit_test.py

@@ -0,0 +1,102 @@
+import mlx.core as mx
+import numpy as np
+
+import keras
+
+
+def test_custom_fit():
+    class CustomModel(keras.Model):
+        def __init__(self, *args, **kwargs):
+            super().__init__(*args, **kwargs)
+            self.loss_tracker = keras.metrics.Mean(name="loss")
+            self.mae_metric = keras.metrics.MeanAbsoluteError(name="mae")
+            self.loss_fn = keras.losses.MeanSquaredError()
+
+        def compute_loss_and_updates(
+            self,
+            trainable_variables,
+            non_trainable_variables,
+            x,
+            y,
+            training=False,
+        ):
+            y_pred, non_trainable_variables = self.stateless_call(
+                trainable_variables,
+                non_trainable_variables,
+                x,
+                training=training,
+            )
+            loss = self.loss_fn(y, y_pred)
+            return loss, (y_pred, non_trainable_variables)
+
+        def train_step(self, state, data):
+            (
+                trainable_variables,
+                non_trainable_variables,
+                optimizer_variables,
+                metrics_variables,
+            ) = state
+            x, y = data
+            grad_fn = mx.value_and_grad(self.compute_loss_and_updates)
+            (loss, (y_pred, non_trainable_variables)), grads = grad_fn(
+                trainable_variables,
+                non_trainable_variables,
+                x,
+                y,
+                training=True,
+            )
+            (
+                trainable_variables,
+                optimizer_variables,
+            ) = self.optimizer.stateless_apply(
+                optimizer_variables, grads, trainable_variables
+            )
+            loss_tracker_vars = metrics_variables[
+                : len(self.loss_tracker.variables)
+            ]
+            mae_metric_vars = metrics_variables[
+                len(self.loss_tracker.variables) :
+            ]
+            loss_tracker_vars = self.loss_tracker.stateless_update_state(
+                loss_tracker_vars, loss
+            )
+            mae_metric_vars = self.mae_metric.stateless_update_state(
+                mae_metric_vars, y, y_pred
+            )
+            logs = {}
+            logs[self.loss_tracker.name] = self.loss_tracker.stateless_result(
+                loss_tracker_vars
+            )
+            logs[self.mae_metric.name] = self.mae_metric.stateless_result(
+                mae_metric_vars
+            )
+            new_metrics_vars = loss_tracker_vars + mae_metric_vars
+            state = (
+                trainable_variables,
+                non_trainable_variables,
+                optimizer_variables,
+                new_metrics_vars,
+            )
+            return logs, state
+
+        @property
+        def metrics(self):
+            return [self.loss_tracker, self.mae_metric]
+
+    inputs = keras.Input(shape=(32,))
+    outputs = keras.layers.Dense(1)(inputs)
+    model = CustomModel(inputs, outputs)
+    model.compile(optimizer="adam")
+    x = np.random.random((64, 32))
+    y = np.random.random((64, 1))
+    history = model.fit(x, y, epochs=1)
+
+    assert "loss" in history.history
+    assert "mae" in history.history
+
+    print("History:")
+    print(history.history)
+
+
+if __name__ == "__main__":
+    test_custom_fit()

keras/src/backend/common/dtypes_test.py

@@ -29,6 +29,12 @@ class DtypesTest(test_case.TestCase):
             for x in dtypes.ALLOWED_DTYPES
             if x not in ["string", "complex64", "complex128"]
         ] + [None]
+    elif backend.backend() == "mlx":
+        ALL_DTYPES = [
+            x
+            for x in dtypes.ALLOWED_DTYPES
+            if x not in ["string", "complex128"]
+        ] + [None]
     else:
         ALL_DTYPES = [x for x in dtypes.ALLOWED_DTYPES if x != "string"] + [
             None

keras/src/backend/mlx/core.py

@@ -25,7 +25,7 @@
 MLX_DTYPES = {
     "float16": mx.float16,
     "float32": mx.float32,
-    "float64": None,  # mlx does not support float64
+    "float64": None,  # mlx only supports float64 on cpu
     "uint8": mx.uint8,
     "uint16": mx.uint16,
     "uint32": mx.uint32,

keras/src/backend/mlx/numpy.py

@@ -1,4 +1,5 @@
 import builtins
+import math
 from copy import copy as builtin_copy
 
 import mlx.core as mx
@@ -950,7 +951,6 @@ def quantile(x, q, axis=None, method="linear", keepdims=False):
     else:
         dtype = dtypes.result_type(x.dtype, float)
     mlx_dtype = to_mlx_dtype(dtype)
-    print("mlx_dtype", mlx_dtype)
 
     # problem casting mlx bfloat16 array to numpy
     if ori_dtype == "bfloat16":
@@ -1374,8 +1374,43 @@ def wrapped(*args):
     return wrapped
 
 
-def histogram(x, bins, range):
-    raise NotImplementedError("histogram not yet implemented in mlx.")
+def histogram_bin_edges(a, bins=10, range=None):
+    # Ref: jax.numpy.histogram
+    # infer range if None
+    if range is None:
+        range = (mx.min(a).item(), mx.max(a).item())
+
+    if range[0] == range[1]:
+        range = (range[0] - 0.5, range[1] + 0.5)
+
+    bin_edges = mx.linspace(range[0], range[1], bins + 1, dtype=mx.float32)
+    # due to the way mlx currently handles linspace
+    # with fp32 precision it is not always right edge inclusive
+    # manually set the right edge for now
+    bin_edges[-1] = range[-1]
+    return bin_edges
+
+
+def histogram(x, bins=10, range=None):
+    # Ref: jax.numpy.histogram
+    x = convert_to_tensor(x)
+    if range is not None:
+        if not isinstance(range, tuple) or len(range) != 2:
+            raise ValueError(
+                "Invalid value for argument `range`. Only `None` or "
+                "a tuple of the lower and upper range of bins is supported. "
+                f"Received: range={range}"
+            )
+
+    bin_edges = histogram_bin_edges(x, bins, range)
+
+    bin_idx = searchsorted(bin_edges, x, side="right")
+    bin_idx = mx.where(x == bin_edges[-1], len(bin_edges) - 1, bin_idx)
+
+    counts = mx.zeros(len(bin_edges))
+    counts = counts.at[bin_idx].add(mx.ones_like(x))
+
+    return counts[1:], bin_edges
 
 
 def unravel_index(x, shape):
@@ -1384,7 +1419,7 @@ def unravel_index(x, shape):
 
     if None in shape:
         raise ValueError(
-            "`shape` argument cannot contain `None`. Received: shape={shape}"
+            f"`shape` argument cannot contain `None`. Received: shape={shape}"
         )
 
     if x.ndim == 1:
@@ -1403,8 +1438,73 @@ def unravel_index(x, shape):
     return tuple(reversed(coords))
 
 
+def searchsorted_binary(a, b, side="left"):
+    original_shape = b.shape
+    b_flat = b.reshape(-1)
+
+    size = a.shape[0]
+    steps = math.ceil(math.log2(size))
+    indices = mx.full(b_flat.shape, vals=size // 2, dtype=mx.uint32)
+
+    comparison = lambda x, y: x <= y if side == "left" else lambda x, y: x < y
+
+    upper = size
+    lower = 0
+    for _ in range(steps):
+        comp = comparison(b_flat, a[indices])
+        new_indices = mx.where(
+            comp, (lower + indices) // 2, (indices + upper) // 2
+        )
+        lower = mx.where(comp, lower, indices)
+        upper = mx.where(comp, indices, upper)
+        indices = new_indices
+
+    result = mx.where(comparison(b_flat, a[indices]), indices, indices + 1)
+    return result.reshape(original_shape)
+
+
+def searchsorted_linear(a, b, side="left"):
+    original_shape = b.shape
+    b_flat = b.reshape(-1)
+    b_flat_broadcast = b_flat.reshape(-1, 1)
+    if side == "left":
+        result = (a[None, :] < b_flat_broadcast).sum(axis=1)
+    else:
+        result = (a[None, :] <= b_flat_broadcast).sum(axis=1)
+
+    return result.reshape(original_shape)
+
+
 def searchsorted(sorted_sequence, values, side="left"):
-    raise NotImplementedError("searchsorted not yet implemented in mlx.")
+    if side not in ("left", "right"):
+        raise ValueError(f"Invalid side `{side}`, must be `left` or `right`.")
+    sorted_sequence = convert_to_tensor(sorted_sequence)
+    values = convert_to_tensor(values)
+    if sorted_sequence.ndim != 1:
+        raise ValueError(
+            "Invalid sorted_sequence, should be 1-dimensional. "
+            f"Recieved sorted_sequence.shape={sorted_sequence.shape}"
+        )
+    if values.ndim == 0:
+        raise ValueError(
+            "Invalid values, should be N-dimensional. Recieved "
+            f"scalar array values.shape={values.shape}"
+        )
+
+    sorted_size = sorted_sequence.size
+    search_size = values.size
+
+    # TODO: swap to mlx implementation if exists in the future
+    # current implementation and search choice based on discussion:
+    # https://github.com/ml-explore/mlx/issues/1255
+    use_linear = sorted_size <= 1024 or (
+        sorted_size <= 16384 and search_size <= 256
+    )
+
+    if use_linear:
+        return searchsorted_linear(sorted_sequence, values, side=side)
+    else:
+        return searchsorted_binary(sorted_sequence, values, side=side)
 
 
 def diagflat(x, k=0):

keras/src/backend/mlx/trainer.py

@@ -260,7 +260,9 @@ def train_step(self, state, data):
                 for ref_v, v in zip(self.metrics_variables, metrics_variables)
             ]
         ) as scope:
-            self._loss_tracker.update_state(unscaled_loss)
+            self._loss_tracker.update_state(
+                unscaled_loss, sample_weight=tree.flatten(x)[0].shape[0]
+            )
             logs = self.compute_metrics(x, y, y_pred, sample_weight)
 
         new_metrics_variables = []
@@ -553,6 +555,7 @@ def fit(
 
         self.stop_training = False
         self.make_train_function()
+        training_logs = {}
         callbacks.on_train_begin()
         initial_epoch = self._initial_epoch or initial_epoch
         for epoch in range(initial_epoch, epochs):
@@ -648,6 +651,7 @@ def fit(
         # If _eval_epoch_iterator exists, delete it after all epochs are done.
         if getattr(self, "_eval_epoch_iterator", None) is not None:
             del self._eval_epoch_iterator
+
         callbacks.on_train_end(logs=training_logs)
         self._mlx_state = None
         return self.history
@@ -706,7 +710,7 @@ def evaluate(
         self.make_test_function()
         self.stop_evaluating = False
         callbacks.on_test_begin()
-        logs = None
+        logs = {}
         self.reset_metrics()
 
         trainable_variables = [v.value for v in self.trainable_variables]