patch for pytest to run with mlx

keras/src/backend/__init__.py

@@ -57,7 +57,8 @@
     distribution_lib = None
 elif backend() == "mlx":
     from keras.src.backend.mlx import *  # noqa: F403
-
+    from keras.src.backend.mlx.core import Variable as BackendVariable
+
     distribution_lib = None
 else:
     raise ValueError(f"Unable to import backend : {backend()}")

keras/src/backend/mlx/__init__.py

@@ -1,5 +1,6 @@
 """MLX backend APIs."""
 
+from keras.src.backend.common.name_scope import name_scope
 from keras.src.backend.mlx import core
 from keras.src.backend.mlx import image
 from keras.src.backend.mlx import linalg

keras/src/backend/mlx/export.py

@@ -0,0 +1,10 @@
+class MlxExportArchive:
+    def track(self, resource):
+        raise NotImplementedError(
+            "`track` is not implemented in the mlx backend."
+        )
+
+    def add_endpoint(self, name, fn, input_signature=None, **kwargs):
+        raise NotImplementedError(
+            "`add_endpoint` is not implemented in the mlx backend."
+        )

keras/src/backend/mlx/nn.py

@@ -115,6 +115,11 @@ def gelu_tanh_approx(x):
     return f(x)
 
 
+def celu(x, alpha=1.0):
+    x = convert_to_tensor(x)
+    return nn.celu(x, alpha=alpha)
+
+
 def softmax(x, axis=-1):
     x = convert_to_tensor(x)
     return mx.softmax(x, axis=axis)

keras/src/export/saved_model.py

@@ -29,6 +29,10 @@
     from keras.src.backend.openvino.export import (
         OpenvinoExportArchive as BackendExportArchive,
     )
+elif backend.backend() == "mlx":
+    from keras.src.backend.mlx.export import (
+        MlxExportArchive as BackendExportArchive,
+    )
 else:
     raise RuntimeError(
         f"Backend '{backend.backend()}' must implement a layer mixin class."