Support diagflat, solve and fix trunc function (#20948)

Author: fbadine

keras/src/backend/mlx/linalg.py

@@ -48,7 +48,10 @@ def lu_factor(a):
 
 
 def solve(a, b):
-    raise NotImplementedError("solve_triangular not yet implemented in mlx.")
+    with mx.stream(mx.cpu):
+        # [linalg::solve] This op is not yet supported on the GPU.
+        # Explicitly pass a CPU stream to run it.
+        return mx.linalg.solve(a, b)
 
 
 def solve_triangular(a, b, lower=False):

keras/src/backend/mlx/numpy.py

@@ -1151,7 +1151,7 @@ def trunc(x):
     dtype = standardize_dtype(x.dtype)
     if "int" in dtype or "bool" == dtype:
         return x
-    return mx.floor(x)
+    return mx.where(x < 0, mx.ceil(x), mx.floor(x))
 
 
 def vdot(x1, x2):
@@ -1514,7 +1514,13 @@ def searchsorted(sorted_sequence, values, side="left"):
 
 
 def diagflat(x, k=0):
-    raise NotImplementedError("diagflat not yet implemented in mlx.")
+    x = convert_to_tensor(x)
+
+    # GPU scatter does not yet support int64 or complex64
+    # for the input or updates.
+    stream = mx.cpu if x.dtype in [mx.int64, mx.complex64] else None
+
+    return mx.diag(mx.reshape(x, [-1]), k, stream=stream)
 
 
 def rot90(array, k=1, axes=(0, 1)):