feat(metal): add generic linalg inverse

Use MPS LU factorization and solve for arbitrary float32 square matrices and batches.

Author: ternaus

CMakeLists.txt

@@ -166,7 +166,9 @@ if(MLX_BUILD_CUDA)
 endif()
 
 if(MLX_BUILD_METAL)
+  enable_language(OBJCXX)
   find_library(METAL_LIB Metal)
+  find_library(METAL_PERFORMANCE_SHADERS_LIB MetalPerformanceShaders)
   find_library(FOUNDATION_LIB Foundation)
   find_library(QUARTZ_LIB QuartzCore)
   if(METAL_LIB)

mlx/backend/metal/CMakeLists.txt

@@ -133,6 +133,7 @@ target_sources(
           ${CMAKE_CURRENT_SOURCE_DIR}/fft.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/hadamard.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/indexing.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/inverse.mm
           ${CMAKE_CURRENT_SOURCE_DIR}/logsumexp.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/matmul.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/scaled_dot_product_attention.cpp
@@ -151,6 +152,8 @@ target_sources(
           ${CMAKE_CURRENT_SOURCE_DIR}/resident.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/utils.cpp)
 
+target_link_libraries(mlx PRIVATE ${METAL_PERFORMANCE_SHADERS_LIB})
+
 if(NOT MLX_METAL_PATH)
   set(MLX_METAL_PATH ${CMAKE_CURRENT_BINARY_DIR}/kernels/)
 endif()

mlx/backend/metal/eval.cpp

@@ -63,7 +63,7 @@ void eval(array& arr) {
       scheduler::notify_task_completion(s);
     });
   } else {
-    command_buffer->addCompletedHandler(
+    encoder.get_command_buffer()->addCompletedHandler(
         [buffers = std::move(buffers)](MTL::CommandBuffer* cbuf) {});
   }
 }

mlx/backend/metal/inverse.mm

@@ -0,0 +1,143 @@
+// Copyright © 2026 Apple Inc.
+
+#include <MetalPerformanceShaders/MetalPerformanceShaders.h>
+
+#include "mlx/allocator.h"
+#include "mlx/backend/metal/device.h"
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+namespace {
+
+id<MTLBuffer> metal_buffer(const array& a) {
+  return (__bridge id<MTLBuffer>)(void*)a.buffer().ptr();
+}
+
+id<MTLDevice> metal_device(metal::Device& device) {
+  return (__bridge id<MTLDevice>)(void*)device.mtl_device();
+}
+
+id<MTLCommandBuffer> metal_command_buffer(metal::CommandEncoder& encoder) {
+  return (__bridge id<MTLCommandBuffer>)(void*)encoder.get_command_buffer();
+}
+
+MPSMatrixDescriptor*
+matrix_descriptor(int rows, int columns, MPSDataType dtype) {
+  const auto row_bytes = static_cast<NSUInteger>(columns) *
+      (dtype == MPSDataTypeUInt32 ? sizeof(uint32_t) : sizeof(float));
+  return [MPSMatrixDescriptor matrixDescriptorWithRows:rows
+                                               columns:columns
+                                              rowBytes:row_bytes
+                                             dataType:dtype];
+}
+
+} // namespace
+
+void Inverse::eval_gpu(const std::vector<array>& inputs, array& output) {
+  if (inputs[0].dtype() != float32) {
+    throw std::invalid_argument(
+        "[Inverse::eval_gpu] Metal inversion supports float32 arrays only.");
+  }
+
+  output.set_data(allocator::malloc(output.nbytes()));
+  if (output.size() == 0) {
+    return;
+  }
+
+  auto& encoder = metal::get_command_encoder(stream());
+  const auto& input = inputs[0];
+  const auto& rhs = inputs[1];
+
+  const int order = input.shape(-1);
+  const size_t batch_size = input.size() / (order * order);
+  array lu(input.shape(), float32, nullptr, {});
+  array pivots({static_cast<int>(batch_size), order}, uint32, nullptr, {});
+  lu.set_data(allocator::malloc(lu.nbytes()));
+  pivots.set_data(allocator::malloc(pivots.nbytes()));
+
+  // MPS needs a separate command buffer after MLX compute work.
+  encoder.end_encoding();
+  auto retained_inputs = std::make_shared<std::vector<array>>(
+      std::initializer_list<array>{input, rhs});
+  auto input_command_buffer = metal_command_buffer(encoder);
+  [input_command_buffer addCompletedHandler:^(id<MTLCommandBuffer>) {
+    (void)retained_inputs;
+  }];
+  encoder.commit();
+
+  auto matrix = matrix_descriptor(order, order, MPSDataTypeFloat32);
+  auto pivot_matrix = matrix_descriptor(1, order, MPSDataTypeUInt32);
+  auto decomposition = [[MPSMatrixDecompositionLU alloc]
+      initWithDevice:metal_device(metal::device(stream().device))
+                rows:order
+             columns:order];
+  auto solve = [[MPSMatrixSolveLU alloc]
+      initWithDevice:metal_device(metal::device(stream().device))
+           transpose:NO
+               order:order
+numberOfRightHandSides:order];
+  auto command_buffer = metal_command_buffer(encoder);
+  auto retained_arrays = std::make_shared<std::vector<array>>(
+      std::initializer_list<array>{input, rhs, lu, pivots});
+  auto resources = [[NSMutableArray alloc] init];
+  [resources addObject:matrix];
+  [resources addObject:pivot_matrix];
+  [resources addObject:decomposition];
+  [resources addObject:solve];
+  [decomposition release];
+  [solve release];
+
+  const auto matrix_bytes = static_cast<size_t>(order) * order * sizeof(float);
+  const auto pivot_bytes = static_cast<size_t>(order) * sizeof(uint32_t);
+  for (size_t batch = 0; batch < batch_size; ++batch) {
+    auto source = [[MPSMatrix alloc] initWithBuffer:metal_buffer(input)
+                                             offset:input.offset() + batch * matrix_bytes
+                                         descriptor:matrix];
+    auto factor = [[MPSMatrix alloc] initWithBuffer:metal_buffer(lu)
+                                             offset:lu.offset() + batch * matrix_bytes
+                                         descriptor:matrix];
+    auto right_hand_side = [[MPSMatrix alloc]
+        initWithBuffer:metal_buffer(rhs)
+                 offset:rhs.offset() + batch * matrix_bytes
+             descriptor:matrix];
+    auto solution = [[MPSMatrix alloc] initWithBuffer:metal_buffer(output)
+                                               offset:output.offset() + batch * matrix_bytes
+                                           descriptor:matrix];
+    auto pivot_indices = [[MPSMatrix alloc]
+        initWithBuffer:metal_buffer(pivots)
+                 offset:pivots.offset() + batch * pivot_bytes
+             descriptor:pivot_matrix];
+    [resources addObject:source];
+    [resources addObject:factor];
+    [resources addObject:right_hand_side];
+    [resources addObject:solution];
+    [resources addObject:pivot_indices];
+    [source release];
+    [factor release];
+    [right_hand_side release];
+    [solution release];
+    [pivot_indices release];
+
+    [decomposition encodeToCommandBuffer:command_buffer
+                             sourceMatrix:source
+                             resultMatrix:factor
+                             pivotIndices:pivot_indices
+                                   status:nil];
+    [solve encodeToCommandBuffer:command_buffer
+                    sourceMatrix:factor
+             rightHandSideMatrix:right_hand_side
+                    pivotIndices:pivot_indices
+                  solutionMatrix:solution];
+  }
+  [command_buffer addCompletedHandler:^(id<MTLCommandBuffer>) {
+    (void)retained_arrays;
+    [resources release];
+  }];
+
+  encoder.register_output_array(output);
+  // Register the MPS result with MLX's command-encoder dependency tracking.
+  encoder.set_input_array(output, 0);
+  encoder.end_encoding();
+}
+
+} // namespace mlx::core

mlx/backend/metal/primitives.cpp

@@ -187,10 +187,6 @@ void SVD::eval_gpu(
   throw std::runtime_error("[SVD::eval_gpu] Metal SVD NYI.");
 }
 
-void Inverse::eval_gpu(const std::vector<array>& inputs, array& output) {
-  throw std::runtime_error("[Inverse::eval_gpu] Metal inversion NYI.");
-}
-
 void Cholesky::eval_gpu(const std::vector<array>& inputs, array& out) {
   throw std::runtime_error(
       "[Cholesky::eval_gpu] Metal Cholesky decomposition NYI.");

mlx/linalg.cpp

@@ -299,7 +299,9 @@ svd(const array& a, bool compute_uv, StreamOrDevice s /* = {} */) {
 }
 
 array inv_impl(const array& a, bool tri, bool upper, StreamOrDevice s) {
-  check_cpu_stream(s, "[linalg::inv]");
+  if (tri) {
+    check_cpu_stream(s, "[linalg::inv]");
+  }
   check_float(a.dtype(), "[linalg::inv]");
 
   if (a.ndim() < 2) {
@@ -314,11 +316,21 @@ array inv_impl(const array& a, bool tri, bool upper, StreamOrDevice s) {
         "[linalg::inv] Inverses are only defined for square matrices.");
   }
 
+  auto stream = to_stream(s);
+  if (stream.device == Device::gpu && !tri) {
+    auto input = contiguous(a, false, stream);
+    auto identity = contiguous(
+        broadcast_to(eye(a.shape(-1), a.dtype(), stream), a.shape(), stream),
+        false,
+        stream);
+    return array(
+        a.shape(),
+        a.dtype(),
+        std::make_shared<Inverse>(stream, tri, upper),
+        {input, identity});
+  }
   return array(
-      a.shape(),
-      a.dtype(),
-      std::make_shared<Inverse>(to_stream(s), tri, upper),
-      {a});
+      a.shape(), a.dtype(), std::make_shared<Inverse>(stream, tri, upper), {a});
 }
 
 array inv(const array& a, StreamOrDevice s /* = {} */) {

python/tests/test_linalg.py

@@ -289,6 +289,19 @@ def test_tri_inverse(self):
         self.assertTrue(mx.all(y1 == mx.triu(y1)))
         self.assertTrue(mx.all(y2 == mx.tril(y2)))
 
+    def test_inverse_generic_matrices(self):
+        rng = np.random.default_rng(0)
+        matrices = rng.normal(size=(2, 7, 7)).astype(np.float32)
+        matrices += 7 * np.eye(7, dtype=np.float32)
+
+        expected = matrices[:, :, ::-1]
+        inverse = mx.linalg.inv(mx.array(matrices)[:, :, ::-1])
+
+        self.assertEqual(inverse.shape, expected.shape)
+        self.assertTrue(
+            np.allclose(inverse, np.linalg.inv(expected), rtol=1e-5, atol=1e-5)
+        )
+
     def test_cholesky(self):
         sqrtA = mx.array(
             [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]], dtype=mx.float32