Add MLX where.

Author: fchollet

keras/backend/mlx/numpy.py

@@ -889,9 +889,7 @@ def vstack(xs):
 
 
 def where(condition, x1, x2):
-    # TODO: Trivial to implement with masking but it would be incorrect for
-    #       instance in the presence of nans or infs
-    raise NotImplementedError("The MLX backend doesn't support where yet")
+    return mx.where(condition, x1, x2)
 
 
 def divide(x1, x2):

keras/backend/mlx/rnn.py

@@ -1,6 +1,3 @@
-import mlx.core as mx
-
-
 def rnn(
     step_function,
     inputs,
@@ -27,7 +24,3 @@ def lstm(*args, **kwargs):
 
 def gru(*args, **kwargs):
     raise NotImplementedError("gru not yet implemented in mlx")
-
-
-def unstack(x, axis=0):
-    return mx.split(x, axis=axis)