Add free threaded build

Co-authored-by: Kanishk <kanishk.chores@gmail.com>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

.github/actions/build-cuda-release/action.yml

@@ -20,7 +20,7 @@ runs:
       run: |
         pip install auditwheel "build<=1.4.2" patchelf setuptools
         python setup.py clean --all
-        MLX_DISABLE_SM90A_KERNELS=1 MLX_BUILD_STAGE=2 python -m build -w
+        MLX_BUILD_STAGE=2 python -m build -w
 
         auditwheel repair dist/mlx_cuda*.whl \
           --plat manylinux_2_35_${{ inputs.arch }} \

.github/actions/build-macos-release/action.yml

@@ -21,7 +21,7 @@ runs:
         DEVELOPER_DIR: /Applications/Xcode-latest.app
         MACOSX_DEPLOYMENT_TARGET: ${{ inputs.macos-target }}
       run: |
-        pip install build
+        uv pip install build
         python setup.py clean --all
         MLX_BUILD_STAGE=1 python -m build -w
 

.github/actions/build-macos/action.yml

@@ -4,61 +4,72 @@ description: 'Build and test MLX on macOS'
 runs:
   using: "composite"
   steps:
-    - name: Install dependencies
+    - name: Install Python package
       env:
         DEBUG: 1
         CMAKE_ARGS: "-DCMAKE_COMPILE_WARNING_AS_ERROR=ON"
-      shell: bash -l {0}
+      shell: bash
       run: |
-        pip install --upgrade pip
-        pip install cmake setuptools typing_extensions
-        pip install -e ".[dev]" -v
+        echo "::group::Install Python package"
+        uv pip install -e ".[dev]" -v
+        echo "::endgroup::"
 
     - name: Install tests dependencies
-      shell: bash -l {0}
+      shell: bash
       run: |
-        pip install tensorflow
+        echo "::group::Install tests dependencies"
+        uv pip install tensorflow
+        echo "::endgroup::"
 
     - name: Run Python tests
-      shell: bash -l {0}
+      shell: bash
       env:
         LOW_MEMORY: 1
       run: |
+        echo "::group::Run Python tests"
         DEVICE=cpu python -m unittest discover -v python/tests
         DEVICE=gpu METAL_DEVICE_WRAPPER_TYPE=1 METAL_DEBUG_ERROR_MODE=0 python -m unittest discover -v python/tests
         mpirun --bind-to none -host localhost:8 -np 8 -x DYLD_LIBRARY_PATH=/opt/homebrew/lib/ python python/tests/mpi_test_distributed.py
         mlx.launch --verbose -n 8 python/tests/ring_test_distributed.py -v 2> >(tee -a stderr.log >&2)
         if $(grep "\[WARN\]" stderr.log); then echo "Distributed ring test failed"; exit 1; fi
-    
+        echo "::endgroup::"
+
     - name: Build example extension
-      shell: bash -l {0}
+      shell: bash
       run: |
+        echo "::group::Build example extension"
         cd examples/extensions
-        pip install -r requirements.txt
-        python setup.py build_ext --inplace
-        python test.py
-    
+        uv pip install -r requirements.txt
+        uv run --no-project setup.py build_ext --inplace
+        uv run --no-project test.py
+        echo "::endgroup::"
+
     - name: Build CPP only
-      shell: bash -l {0}
+      shell: bash
       run: |
+        echo "::group::Build CPP only"
         mkdir -p build
         cd build
         cmake ..
         make -j $(sysctl -n hw.ncpu)
+        echo "::endgroup::"
 
     - name: Run CPP tests
-      shell: bash -l {0}
+      shell: bash
       env:
         DEVICE: gpu
         METAL_DEVICE_WRAPPER_TYPE: 1
         METAL_DEBUG_ERROR_MODE: 0
       run: |
+        echo "::group::Run CPP tests"
         ./build/tests/tests
         ./build/tests/test_teardown
+        echo "::endgroup::"
 
     - name: Build small binary with JIT
-      shell: bash -l {0}
+      shell: bash
       run: |
+        echo "::group::Build small binary with JIT"
         mkdir -p build
         cd build
         cmake .. -DCMAKE_BUILD_TYPE=MinSizeRel \
@@ -68,15 +79,18 @@ runs:
           -DMLX_BUILD_GGUF=OFF \
           -DMLX_METAL_JIT=ON
         make -j $(sysctl -n hw.ncpu)
-    
+        echo "::endgroup::"
+
     - name: Run Python tests with JIT
-      shell: bash -l {0}
+      shell: bash
       env:
         LOW_MEMORY: 1
         DEVICE: gpu
         METAL_DEVICE_WRAPPER_TYPE: 1
         METAL_DEBUG_ERROR_MODE: 0
       run: |
+        echo "::group::Run Python tests with JIT"
         CMAKE_ARGS="-DMLX_METAL_JIT=ON" \
-          pip install -e . -v
+          uv pip install -e . -v
         python -m unittest discover -v python/tests
+        echo "::endgroup::"

.github/actions/setup-linux/action.yml

@@ -14,6 +14,9 @@ inputs:
     description: 'Whether to enable ccache'
     required: false
     default: 'true'
+  ccache-key:
+    required: false
+    default: 'ccache'
 
 runs:
   using: "composite"
@@ -33,7 +36,7 @@ runs:
       if: ${{ inputs.use-ccache == 'true' }}
       uses: hendrikmuhs/ccache-action@v1.2
       with:
-        key: ccache-${{ runner.os }}-${{ runner.arch }}-${{ inputs.toolkit }}
+        key: ${{ inputs.ccache-key }}-${{ runner.os }}-${{ runner.arch }}-${{ inputs.toolkit }}
         max-size: 1GB
         # ccache-action bug: running "apt-get update" fails on large arm runner.
         update-package-index: false
@@ -55,7 +58,7 @@ runs:
         echo "::endgroup::"
 
     - name: Set swap space
-      if: ${{ startsWith(inputs.toolkit, 'cuda') && runner.arch == 'arm64' }}
+      if: ${{ startsWith(inputs.toolkit, 'cuda') }}
       uses: pierotofy/set-swap-space@fc79b3f67fa8a838184ce84a674ca12238d2c761
       with:
         swap-size-gb: 16

.github/actions/setup-macos/action.yml

@@ -13,12 +13,20 @@ runs:
     - name: Install Homebrew packages
       shell: sh
       run: /opt/homebrew/bin/brew install openmpi
-    
+
     - name: Verify MetalToolchain installed
       shell: bash
       run: xcodebuild -showComponent MetalToolchain
 
-    - uses: conda-incubator/setup-miniconda@v3
-      with:
-        miniconda-version: "latest"
-        python-version: ${{ inputs.python-version }}
+    - uses: astral-sh/setup-uv@v7
+
+    - name: Setup Python venv
+      shell: bash
+      run: |
+        echo "::group::Setup Python venv"
+        uv venv --python ${{ inputs.python-version }}
+        source .venv/bin/activate
+        echo PATH=$PATH >> $GITHUB_ENV
+        # Search python packages in .venv
+        echo PYTHONPATH=`python -c 'import sys; print(sys.path[-1])'` >> $GITHUB_ENV
+        echo "::endgroup::"

.github/actions/test-linux/action.yml

@@ -9,7 +9,21 @@ inputs:
 runs:
   using: "composite"
   steps:
+    # FIXME: The distributed tests fail with free-threading Python.
+    - name: Check free-threading Python
+      id: is-free-threading
+      shell: bash
+      run: |
+        echo "::group::Check free-threading Python"
+        if python -VV 2>&1 | grep "free-threading"; then
+          echo "result=true" >> $GITHUB_OUTPUT
+        else
+          echo "result=false" >> $GITHUB_OUTPUT
+        fi
+        echo "::endgroup::"
+
     - name: Run MPI tests
+      if: ${{ steps.is-free-threading.outputs.result == 'false' }}
       shell: bash
       run: |
         echo "::group::MPI tests"
@@ -17,7 +31,7 @@ runs:
         echo "::endgroup::"
 
     - name: Run distributed tests
-      if: ${{ inputs.has-gpu == 'false' }}
+      if: ${{ steps.is-free-threading.outputs.result == 'false' && inputs.has-gpu == 'false' }}
       shell: bash
       run: |
         echo "::group::Distributed tests"

.github/workflows/nightly.yml

@@ -41,7 +41,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python_version: ["3.11", "3.12", "3.13", "3.14"]
+        python_version: ["3.11", "3.12", "3.13", "3.14", "3.14t"]
         runner:
           - ubuntu-22.04
           - ubuntu-22.04-arm
@@ -59,7 +59,7 @@ jobs:
     if: github.repository == 'ml-explore/mlx'
     strategy:
       matrix:
-        python-version: ["3.10", "3.13"]
+        python-version: ["3.10", "3.13", "3.14t"]
     runs-on: [self-hosted, macos]
     steps:
       - uses: actions/checkout@v6
@@ -85,20 +85,24 @@ jobs:
 
   build_cuda_release:
     if: github.repository == 'ml-explore/mlx'
-    runs-on: ubuntu-22-large
+    strategy:
+      matrix:
+        arch: ['x86_64', 'aarch64']
+        toolkit: ['cuda-12.9', 'cuda-13.0']
+    runs-on: ${{ matrix.arch == 'x86_64' && 'ubuntu-22-large' || 'ubuntu-22-large-arm' }}
     steps:
       - uses: actions/checkout@v6
       - uses: ./.github/actions/setup-linux
         with:
-          toolkit: 'cuda-12.9'
+          toolkit: ${{ matrix.toolkit }}
+          ccache-key: 'ccache-release'
       - name: Build Python package
         uses: ./.github/actions/build-cuda-release
         with:
-          toolkit: 'cuda-12.9'
-          arch: 'x86_64'
+          arch: ${{ matrix.arch }}
       - name: Upload artifacts
         uses: actions/upload-artifact@v7
         with:
-          name: mlx-cuda
+          name: mlx-${{ matrix.toolkit }}-${{ matrix.arch }}
           path: wheelhouse/mlx_cuda_*.whl
           retention-days: 7

.github/workflows/release.yml

@@ -47,7 +47,7 @@ jobs:
     if: github.repository == 'ml-explore/mlx'
     strategy:
       matrix:
-        python_version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
+        python_version: ["3.10", "3.11", "3.12", "3.13", "3.13t", "3.14", "3.14t"]
         arch: ['x86_64', 'aarch64']
     runs-on: ${{ matrix.arch == 'x86_64' && 'ubuntu-22.04' || 'ubuntu-22.04-arm' }}
     env:
@@ -83,7 +83,7 @@ jobs:
     if: github.repository == 'ml-explore/mlx'
     strategy:
       matrix:
-        python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
+        python-version: ["3.10", "3.11", "3.12", "3.13", "3.13t", "3.14", "3.14t"]
     runs-on: [self-hosted, macos]
     env:
       PYPI_RELEASE: 1
@@ -93,13 +93,8 @@ jobs:
       - uses: ./.github/actions/setup-macos
         with:
           python-version: ${{ matrix.python-version }}
-
-      - name: Install dependencies
-        shell: bash -l {0}
-        run: |
-          pip install --upgrade pip
-          pip install cmake setuptools typing_extensions
-          pip install -e . -v
+      - name: Install Python package
+        run: uv pip install -e . -v
       - name: Build macOS 14 package
         uses: ./.github/actions/build-macos-release
         with:
@@ -146,7 +141,7 @@ jobs:
       - uses: ./.github/actions/setup-linux
         with:
           toolkit: ${{ matrix.toolkit }}
-          use-ccache: false
+          ccache-key: 'ccache-release'
       - name: Build Python package
         uses: ./.github/actions/build-cuda-release
         with:

docs/src/python/linalg.rst

@@ -14,6 +14,7 @@ Linear Algebra
     cholesky
     cholesky_inv
     cross
+    det
     qr
     svd
     eigvals
@@ -23,5 +24,6 @@ Linear Algebra
     lu
     lu_factor
     pinv
+    slogdet
     solve
     solve_triangular

examples/extensions/requirements.txt

@@ -1,4 +1,4 @@
 setuptools>=42
 cmake>=3.25
-mlx>=0.21.0
+mlx>=0.31.2
 nanobind==2.12.0

mlx/backend/cpu/luf.cpp

@@ -67,11 +67,10 @@ void luf_impl(
               /* ipiv */ reinterpret_cast<int*>(pivots_ptr),
               /* info */ &info);
 
-          if (info != 0) {
+          if (info < 0) {
             std::stringstream ss;
             ss << "[LUF::eval_cpu] sgetrf_ failed with code " << info
-               << ((info > 0) ? " because matrix is singular"
-                              : " because argument had an illegal value");
+               << " because argument had an illegal value";
             throw std::runtime_error(ss.str());
           }
 

mlx/backend/cuda/CMakeLists.txt

@@ -168,9 +168,8 @@ set_target_properties(mlx PROPERTIES CUDA_ARCHITECTURES
                                      "${MLX_CUDA_ARCHITECTURES}")
 
 # Skip Hopper-only kernels when not building for sm90a.
-if(NOT DEFINED ENV{MLX_DISABLE_SM90A_KERNELS}
-   AND (("90a" IN_LIST MLX_CUDA_ARCHITECTURES) OR ("90a-real" IN_LIST
-                                                   MLX_CUDA_ARCHITECTURES)))
+if(("90a" IN_LIST MLX_CUDA_ARCHITECTURES) OR ("90a-real" IN_LIST
+                                              MLX_CUDA_ARCHITECTURES))
   target_compile_definitions(mlx PRIVATE MLX_CUDA_SM90A_ENABLED)
 endif()
 

mlx/backend/cuda/gemms/gather_gemm.cu

@@ -43,6 +43,12 @@ class GatherGemm {
   using ElementD = typename CollectiveEpilogue::ElementD;
   using StrideD = typename CollectiveEpilogue::StrideD;
 
+  static_assert(
+      cute::is_same_v<
+          ElementAccumulator,
+          typename CollectiveEpilogue::ElementAccumulator>,
+      "Mainloop and epilogue do not agree on accumulator value type.");
+
   static constexpr int SharedStorageSize = static_cast<int>(cute::max(
       sizeof(typename CollectiveMainloop::SharedStorage),
       sizeof(typename CollectiveEpilogue::SharedStorage)));
@@ -98,7 +104,9 @@ class GatherGemm {
 
   CUTLASS_DEVICE void operator()(const Params& params, char* smem_buf) {
     int thread_idx = int(threadIdx.x);
-    auto [m_coord, n_coord, l_coord] = uint3(blockIdx);
+    int m_coord = int(blockIdx.x);
+    int n_coord = int(blockIdx.y);
+    int l_coord = int(blockIdx.z);
 
     auto shape_MNKL = append<4>(params.problem_shape, Int<1>{});
     auto cta_tile = TileShape{};
@@ -220,7 +228,7 @@ void gather_mm(
       using TileShape = Shape<_128, _128, _8>;
       using DispatchPolicy = cutlass::gemm::MainloopSm70TwoStage;
       using TiledMma = TiledMMA<
-          MMA_Atom<UniversalFMA<Accumulator, Element, Element, Element>>,
+          MMA_Atom<UniversalFMA<Accumulator, Element, Element, Accumulator>>,
           Layout<Shape<_16, _16, _1>>>;
 
       using CopyTraitsA = SimtCopyTraits<Element, k_major_a.value>;
@@ -296,9 +304,6 @@ void cutlass_gather_mm(
   int n = out.shape(-1);
   int k = a.shape(-1);
   int l = out.size() / (m * n);
-  if (m < 16 || n < 16) {
-    throw std::invalid_argument("[gather_mm] M/N is too small.");
-  }
 
   encoder.set_input_array(a);
   encoder.set_input_array(b);

mlx/backend/cuda/gemms/grouped_gemm_unaligned.cu

@@ -245,7 +245,7 @@ void grouped_gemm_v2(
           LayoutB,
           cutlass::ComplexTransform::kNone,
           GemmConfiguration::kAlignmentAB,
-          typename GemmConfiguration::Element,
+          typename GemmConfiguration::Accumulator,
           cutlass::layout::RowMajor,
           typename GemmConfiguration::Accumulator,
           typename GemmConfiguration::OpClass,

mlx/backend/cuda/quantized/qmm/CMakeLists.txt

@@ -1,19 +1,35 @@
 target_sources(
   mlx
-  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/qmm.cu
-          ${CMAKE_CURRENT_SOURCE_DIR}/qmv.cu
-          ${CMAKE_CURRENT_SOURCE_DIR}/fp_qmv.cu
-          ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_naive_m16_k.cu
-          ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_naive_m16_n.cu
-          ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_naive_m32_k.cu
-          ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_naive_m32_n.cu
-          ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_naive_m64_k.cu
-          ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_naive_m64_n.cu
-          ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_sm80_m16.cu
-          ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_sm80_m32.cu
-          ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_sm80_m64.cu
-          ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_sm90_m128_n16_m1.cu
-          ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_sm90_m128_n32_m1.cu
-          ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_sm90_m128_n64_m2.cu
-          ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_sm90_m128_n128_m2.cu
-          ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_sm90_m128_n256_m2.cu)
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/qmm.cu ${CMAKE_CURRENT_SOURCE_DIR}/qmv.cu
+          ${CMAKE_CURRENT_SOURCE_DIR}/fp_qmv.cu)
+
+foreach(TileN 16 32 64 128 256)
+  set(OUTPUT_FILE "qmm_sm90_impl_n${TileN}.cu")
+  configure_file("${CMAKE_CURRENT_SOURCE_DIR}/qmm_sm90.cu"
+                 "${CMAKE_CURRENT_BINARY_DIR}/${OUTPUT_FILE}" @ONLY)
+  target_sources(mlx PRIVATE ${CMAKE_CURRENT_BINARY_DIR}/${OUTPUT_FILE})
+endforeach()
+
+foreach(TileM 16 32 64)
+  set(OUTPUT_FILE "qmm_sm80_impl_m${TileM}.cu")
+  configure_file("${CMAKE_CURRENT_SOURCE_DIR}/qmm_sm80.cu"
+                 "${CMAKE_CURRENT_BINARY_DIR}/${OUTPUT_FILE}" @ONLY)
+  target_sources(mlx PRIVATE ${CMAKE_CURRENT_BINARY_DIR}/${OUTPUT_FILE})
+endforeach()
+
+foreach(TileM 16 32 64)
+  foreach(KMajor true false)
+    foreach(HasKResidue true false)
+      foreach(SM80 true false)
+        if(${KMajor} AND ${HasKResidue})
+          continue()
+        endif()
+        set(OUTPUT_FILE
+            "qmm_naive_impl_m${TileM}_${KMajor}_${HasKResidue}_${SM80}.cu")
+        configure_file("${CMAKE_CURRENT_SOURCE_DIR}/qmm_naive.cu"
+                       "${CMAKE_CURRENT_BINARY_DIR}/${OUTPUT_FILE}" @ONLY)
+        target_sources(mlx PRIVATE ${CMAKE_CURRENT_BINARY_DIR}/${OUTPUT_FILE})
+      endforeach()
+    endforeach()
+  endforeach()
+endforeach()

mlx/backend/cuda/quantized/qmm/qmm.cu

@@ -17,9 +17,9 @@ inline bool is_last_2_dims_row_contiguous(const array& x) {
 } // namespace
 
 #if defined(MLX_CUDA_SM90A_ENABLED)
-// Defined in qmm_impl_sm90_xxx.cu files.
-template <typename TileShape, typename ClusterShape>
-void qmm_impl_sm90(
+// Defined in qmm_sm90.cu.
+template <int TileN>
+void qmm_sm90_impl(
     const array& x,
     const array& w,
     const array& scales,
@@ -52,7 +52,7 @@ bool supports_qmm_sm90(
   if (!biases) {
     return false;
   }
-  if (!x.flags().row_contiguous || !is_last_2_dims_row_contiguous(w) ||
+  if (!is_last_2_dims_row_contiguous(w) ||
       !is_last_2_dims_row_contiguous(scales) ||
       !is_last_2_dims_row_contiguous(*biases)) {
     return false;
@@ -83,24 +83,21 @@ void qmm_sm90(
     cu::CommandEncoder& encoder,
     Stream s) {
 #if defined(MLX_CUDA_SM90A_ENABLED)
-  auto dispatch = [&]<int tile_m, int tile_n, int cluster_m>() {
-    using cute::Int;
-    using TileShapeMN = cute::Shape<Int<tile_m>, Int<tile_n>>;
-    using ClusterShape = cute::Shape<Int<cluster_m>, Int<1>, Int<1>>;
-    qmm_impl_sm90<TileShapeMN, ClusterShape>(
+  auto dispatch = [&]<int TileN>() {
+    qmm_sm90_impl<TileN>(
         x, w, scales, biases, out, bits, group_size, encoder, s);
   };
   int m = out.ndim() > 1 ? out.shape(-2) : 1;
   if (m <= 16) {
-    dispatch.template operator()<128, 16, 1>();
+    dispatch.template operator()<16>();
   } else if (m <= 32) {
-    dispatch.template operator()<128, 32, 1>();
+    dispatch.template operator()<32>();
   } else if (m <= 64) {
-    dispatch.template operator()<128, 64, 2>();
+    dispatch.template operator()<64>();
   } else if (m <= 128) {
-    dispatch.template operator()<128, 128, 2>();
+    dispatch.template operator()<128>();
   } else {
-    dispatch.template operator()<128, 256, 2>();
+    dispatch.template operator()<256>();
   }
 #else
   throw std::runtime_error(
@@ -108,9 +105,9 @@ void qmm_sm90(
 #endif // defined(MLX_CUDA_SM90A_ENABLED)
 }
 
-// Defined in qmm_impl_sm80_xxx.cu files.
+// Defined in qmm_sm80.cu.
 template <int TileM>
-void qmm_impl_sm80(
+void qmm_sm80_impl(
     const array& x,
     const array& w,
     const array& scales,
@@ -142,7 +139,7 @@ bool supports_qmm_sm80(
   if ((n % 128 != 0) || (k % std::max(64, group_size) != 0)) {
     return false;
   }
-  if (!x.flags().row_contiguous || !is_last_2_dims_row_contiguous(w) ||
+  if (!is_last_2_dims_row_contiguous(w) ||
       !is_last_2_dims_row_contiguous(scales)) {
     return false;
   }
@@ -174,7 +171,7 @@ void qmm_sm80(
     QuantizationMode mode,
     cu::CommandEncoder& encoder) {
   auto dispatch = [&]<int TileM>() {
-    qmm_impl_sm80<TileM>(
+    qmm_sm80_impl<TileM>(
         x,
         w,
         scales,
@@ -197,9 +194,9 @@ void qmm_sm80(
   }
 }
 
-// Defined in qmm_impl_naive_xxx.cu files.
-template <int TileM, bool KMajor>
-void qmm_impl_naive(
+// Defined in qmm_naive.cu.
+template <int TileM, bool KMajor, bool HasKResidue, bool SM80>
+void qmm_naive_impl(
     const array& x,
     const array& w,
     const array& scales,
@@ -227,7 +224,7 @@ bool supports_qmm_naive(
   if (transpose && (k % std::max(64, group_size) != 0)) {
     return false;
   }
-  if (!x.flags().row_contiguous || !is_last_2_dims_row_contiguous(w) ||
+  if (!is_last_2_dims_row_contiguous(w) ||
       !is_last_2_dims_row_contiguous(scales)) {
     return false;
   }
@@ -250,8 +247,8 @@ void qmm_naive(
     int group_size,
     QuantizationMode mode,
     cu::CommandEncoder& encoder) {
-  auto dispatch = [&]<int TileM, bool KMajor>() {
-    qmm_impl_naive<TileM, KMajor>(
+  auto dispatch = [&]<int TileM, bool KMajor, bool HasKResidue, bool SM80>() {
+    qmm_naive_impl<TileM, KMajor, HasKResidue, SM80>(
         x,
         w,
         scales,
@@ -264,15 +261,38 @@ void qmm_naive(
         mode,
         encoder);
   };
-  dispatch_bool(transpose, [&](auto k_major) {
-    int m = out.ndim() > 1 ? out.shape(-2) : 1;
-    if (m <= 16) {
-      dispatch.template operator()<16, k_major.value>();
-    } else if (m <= 32) {
-      dispatch.template operator()<32, k_major.value>();
+  auto dispatch_k = [&](auto k_major, bool has_k_residue, auto&& f) {
+    if constexpr (k_major.value) {
+      if (has_k_residue) {
+        throw std::invalid_argument(
+            "[quantized_matmul] K must be multiples of max(64, group_size).");
+      }
+      f.template operator()<false>();
     } else {
-      dispatch.template operator()<64, k_major.value>();
+      dispatch_bool(has_k_residue, [&](auto has_k_residue) {
+        f.template operator()<has_k_residue.value>();
+      });
     }
+  };
+  int m = out.ndim() > 1 ? out.shape(-2) : 1;
+  int k = x.shape(-1);
+  int tile_k = std::max(64, group_size);
+  bool has_k_residue = k % tile_k != 0;
+  bool sm80 = encoder.device().compute_capability_major() >= 8;
+  dispatch_bool(transpose, [&](auto k_major) {
+    dispatch_k(k_major, has_k_residue, [&]<bool HasKResidue>() {
+      dispatch_bool(sm80, [&](auto sm80) {
+        constexpr bool KMajor = k_major.value;
+        constexpr bool SM80 = sm80.value;
+        if (m <= 16) {
+          dispatch.template operator()<16, KMajor, HasKResidue, SM80>();
+        } else if (m <= 32) {
+          dispatch.template operator()<32, KMajor, HasKResidue, SM80>();
+        } else {
+          dispatch.template operator()<64, KMajor, HasKResidue, SM80>();
+        }
+      });
+    });
   });
 }
 
@@ -323,7 +343,7 @@ bool supports_qmv(
   if (k % 8 != 0) {
     return false;
   }
-  if (!x.flags().row_contiguous || !is_last_2_dims_row_contiguous(w) ||
+  if (!is_last_2_dims_row_contiguous(w) ||
       !is_last_2_dims_row_contiguous(scales)) {
     return false;
   }

mlx/backend/cuda/quantized/qmm/qmm_impl_naive_m16_k.cu

@@ -1,5 +0,0 @@
-// Copyright © 2026 Apple Inc.
-
-#include "mlx/backend/cuda/quantized/qmm/qmm_impl_naive.cuh"
-
-QMM_NAIVE_GPU(16, true)

mlx/backend/cuda/quantized/qmm/qmm_impl_naive_m16_n.cu

@@ -1,5 +0,0 @@
-// Copyright © 2026 Apple Inc.
-
-#include "mlx/backend/cuda/quantized/qmm/qmm_impl_naive.cuh"
-
-QMM_NAIVE_GPU(16, false)

mlx/backend/cuda/quantized/qmm/qmm_impl_naive_m32_k.cu

@@ -1,5 +0,0 @@
-// Copyright © 2026 Apple Inc.
-
-#include "mlx/backend/cuda/quantized/qmm/qmm_impl_naive.cuh"
-
-QMM_NAIVE_GPU(32, true)

mlx/backend/cuda/quantized/qmm/qmm_impl_naive_m32_n.cu

@@ -1,5 +0,0 @@
-// Copyright © 2026 Apple Inc.
-
-#include "mlx/backend/cuda/quantized/qmm/qmm_impl_naive.cuh"
-
-QMM_NAIVE_GPU(32, false)

mlx/backend/cuda/quantized/qmm/qmm_impl_naive_m64_k.cu

@@ -1,5 +0,0 @@
-// Copyright © 2026 Apple Inc.
-
-#include "mlx/backend/cuda/quantized/qmm/qmm_impl_naive.cuh"
-
-QMM_NAIVE_GPU(64, true)

mlx/backend/cuda/quantized/qmm/qmm_impl_naive_m64_n.cu

@@ -1,5 +0,0 @@
-// Copyright © 2026 Apple Inc.
-
-#include "mlx/backend/cuda/quantized/qmm/qmm_impl_naive.cuh"
-
-QMM_NAIVE_GPU(64, false)

mlx/backend/cuda/quantized/qmm/qmm_impl_sm80_m16.cu

@@ -1,5 +0,0 @@
-// Copyright © 2026 Apple Inc.
-
-#include "mlx/backend/cuda/quantized/qmm/qmm_impl_sm80.cuh"
-
-QMM_SM80_GPU(16)

mlx/backend/cuda/quantized/qmm/qmm_impl_sm80_m32.cu

@@ -1,5 +0,0 @@
-// Copyright © 2026 Apple Inc.
-
-#include "mlx/backend/cuda/quantized/qmm/qmm_impl_sm80.cuh"
-
-QMM_SM80_GPU(32)

mlx/backend/cuda/quantized/qmm/qmm_impl_sm80_m64.cu

@@ -1,5 +0,0 @@
-// Copyright © 2026 Apple Inc.
-
-#include "mlx/backend/cuda/quantized/qmm/qmm_impl_sm80.cuh"
-
-QMM_SM80_GPU(64)

mlx/backend/cuda/quantized/qmm/qmm_impl_sm90_m128_n128_m2.cu

@@ -1,10 +0,0 @@
-// Copyright © 2026 Apple Inc.
-
-#include "mlx/backend/cuda/quantized/qmm/qmm_impl_sm90.cuh"
-
-using namespace cute;
-
-using TileShapeMN = Shape<_128, _128>;
-using ClusterShape = Shape<_2, _1, _1>;
-
-QMM_SM90_GPU(TileShapeMN, ClusterShape)

mlx/backend/cuda/quantized/qmm/qmm_impl_sm90_m128_n16_m1.cu

@@ -1,10 +0,0 @@
-// Copyright © 2026 Apple Inc.
-
-#include "mlx/backend/cuda/quantized/qmm/qmm_impl_sm90.cuh"
-
-using namespace cute;
-
-using TileShapeMN = Shape<_128, _16>;
-using ClusterShape = Shape<_1, _1, _1>;
-
-QMM_SM90_GPU(TileShapeMN, ClusterShape)

mlx/backend/cuda/quantized/qmm/qmm_impl_sm90_m128_n256_m2.cu

@@ -1,10 +0,0 @@
-// Copyright © 2026 Apple Inc.
-
-#include "mlx/backend/cuda/quantized/qmm/qmm_impl_sm90.cuh"
-
-using namespace cute;
-
-using TileShapeMN = Shape<_128, _256>;
-using ClusterShape = Shape<_2, _1, _1>;
-
-QMM_SM90_GPU(TileShapeMN, ClusterShape)

mlx/backend/cuda/quantized/qmm/qmm_impl_sm90_m128_n32_m1.cu

@@ -1,10 +0,0 @@
-// Copyright © 2026 Apple Inc.
-
-#include "mlx/backend/cuda/quantized/qmm/qmm_impl_sm90.cuh"
-
-using namespace cute;
-
-using TileShapeMN = Shape<_128, _32>;
-using ClusterShape = Shape<_1, _1, _1>;
-
-QMM_SM90_GPU(TileShapeMN, ClusterShape)

mlx/backend/cuda/quantized/qmm/qmm_impl_sm90_m128_n64_m2.cu

@@ -1,10 +0,0 @@
-// Copyright © 2026 Apple Inc.
-
-#include "mlx/backend/cuda/quantized/qmm/qmm_impl_sm90.cuh"
-
-using namespace cute;
-
-using TileShapeMN = Shape<_128, _64>;
-using ClusterShape = Shape<_2, _1, _1>;
-
-QMM_SM90_GPU(TileShapeMN, ClusterShape)

mlx/backend/cuda/quantized/qmm/qmm_naive.cu

@@ -60,7 +60,9 @@ __global__ void qmm_naive_kernel(
   CUTE_STATIC_ASSERT_V(congruent(select<0,1,3>(shape_MNKL), dC));
 
   int thread_idx = int(threadIdx.x);
-  auto [m_coord, n_coord, l_coord] = static_cast<uint3>(blockIdx);
+  int m_coord = int(blockIdx.x);
+  int n_coord = int(blockIdx.y);
+  int l_coord = int(blockIdx.z);
 
   auto m_max_coord = size<0>(shape_MNKL) - size<0>(cta_tiler) * m_coord; // M - BLK_M * m_coord
   auto n_max_coord = size<1>(shape_MNKL) - size<1>(cta_tiler) * n_coord; // N - BLK_N * n_coord
@@ -316,7 +318,7 @@ inline constexpr auto make_scales_layout(auto n, auto k, auto l, auto group_size
   }
 }
 
-template <int TileM = 16, bool KMajor = true, bool SM80 = true, bool HasKResidue = false,
+template <int TileM = 16, bool KMajor = true, bool HasKResidue = false, bool SM80 = true,
           typename Element, typename Quant, typename Scale>
 void qmm_naive(
     const Element* A,
@@ -396,21 +398,6 @@ void qmm_naive(
 
 namespace mlx::core {
 
-template <bool KMajor, typename F>
-inline void dispatch_k(bool has_k_residue, const char* tag, F&& f) {
-  if constexpr (KMajor) {
-    if (has_k_residue) {
-      throw std::invalid_argument(
-          fmt::format("{} K must be multiples of group_size.", tag));
-    }
-    f.template operator()<false>();
-  } else {
-    dispatch_bool(has_k_residue, [&](auto has_k_residue) {
-      f.template operator()<has_k_residue.value>();
-    });
-  }
-}
-
 template <typename F>
 inline void dispatch_element_types(Dtype dtype, const char* tag, F&& f) {
   if (dtype == float32) {
@@ -474,8 +461,8 @@ inline void dispatch_quant_types(
   }
 }
 
-template <int TileM, bool KMajor>
-void qmm_impl_naive(
+template <int TileM, bool KMajor, bool HasKResidue, bool SM80>
+void qmm_naive_impl(
     const array& x,
     const array& w,
     const array& scales,
@@ -494,71 +481,65 @@ void qmm_impl_naive(
   int l = out.size() / (m * n);
   bool broadcast_b = (w.ndim() <= 2) || (w.size() != w.data_size());
 
-  bool is_sm80 = encoder.device().compute_capability_major() >= 8;
-  dispatch_bool(is_sm80, [&](auto sm80) {
-    dispatch_k<KMajor>(k % group_size != 0, tag, [&]<bool has_k_residue>() {
-      dispatch_element_types(out.dtype(), tag, [&]<typename Element>() {
-        dispatch_quant_types<Element>(
-            bits,
-            group_size,
-            mode,
-            tag,
-            [&]<typename Quant, typename Scale, int group_size>() {
-              encoder.set_input_array(x);
-              encoder.set_input_array(w);
-              encoder.set_input_array(scales);
-              if (biases) {
-                encoder.set_input_array(*biases);
-              }
-              if (lhs_indices) {
-                encoder.set_input_array(*lhs_indices);
-              }
-              if (rhs_indices) {
-                encoder.set_input_array(*rhs_indices);
-              }
-              encoder.set_output_array(out);
-              cutlass_gemm::qmm_naive<TileM, KMajor, sm80.value, has_k_residue>(
-                  gpu_ptr<Element>(x),
-                  gpu_ptr<Quant>(w),
-                  gpu_ptr<Scale>(scales),
-                  biases ? gpu_ptr<Element>(*biases) : nullptr,
-                  lhs_indices ? gpu_ptr<uint32_t>(*lhs_indices) : nullptr,
-                  rhs_indices ? gpu_ptr<uint32_t>(*rhs_indices) : nullptr,
-                  gpu_ptr<Element>(out),
-                  m,
-                  n,
-                  k,
-                  l,
-                  broadcast_b,
-                  cute::Int<group_size>{},
-                  [&](auto* kernel,
-                      dim3 num_blocks,
-                      dim3 block_dims,
-                      uint32_t smem_bytes,
-                      void** args) {
-                    encoder.add_kernel_node_raw(
-                        kernel, num_blocks, block_dims, {}, smem_bytes, args);
-                  });
-            });
-      });
-    });
+  dispatch_element_types(out.dtype(), tag, [&]<typename Element>() {
+    dispatch_quant_types<Element>(
+        bits,
+        group_size,
+        mode,
+        tag,
+        [&]<typename Quant, typename Scale, int group_size>() {
+          encoder.set_input_array(x);
+          encoder.set_input_array(w);
+          encoder.set_input_array(scales);
+          if (biases) {
+            encoder.set_input_array(*biases);
+          }
+          if (lhs_indices) {
+            encoder.set_input_array(*lhs_indices);
+          }
+          if (rhs_indices) {
+            encoder.set_input_array(*rhs_indices);
+          }
+          encoder.set_output_array(out);
+          cutlass_gemm::qmm_naive<TileM, KMajor, HasKResidue, SM80>(
+              gpu_ptr<Element>(x),
+              gpu_ptr<Quant>(w),
+              gpu_ptr<Scale>(scales),
+              biases ? gpu_ptr<Element>(*biases) : nullptr,
+              lhs_indices ? gpu_ptr<uint32_t>(*lhs_indices) : nullptr,
+              rhs_indices ? gpu_ptr<uint32_t>(*rhs_indices) : nullptr,
+              gpu_ptr<Element>(out),
+              m,
+              n,
+              k,
+              l,
+              broadcast_b,
+              cute::Int<group_size>{},
+              [&](auto* kernel,
+                  dim3 num_blocks,
+                  dim3 block_dims,
+                  uint32_t smem_bytes,
+                  void** args) {
+                encoder.add_kernel_node_raw(
+                    kernel, num_blocks, block_dims, {}, smem_bytes, args);
+              });
+        });
   });
 }
 
-} // namespace mlx::core
+// clang-format off
+template void qmm_naive_impl<@TileM@, @KMajor@, @HasKResidue@, @SM80@>(
+    const array& x,
+    const array& w,
+    const array& scales,
+    const std::optional<array>& biases,
+    const std::optional<array>& lhs_indices,
+    const std::optional<array>& rhs_indices,
+    array& out,
+    int bits,
+    int group_size,
+    QuantizationMode mode,
+    cu::CommandEncoder& encoder);
+// clang-format on
 
-#define QMM_NAIVE_GPU(TileM, KMajor)           \
-  namespace mlx::core {                        \
-  template void qmm_impl_naive<TileM, KMajor>( \
-      const array& x,                          \
-      const array& w,                          \
-      const array& scales,                     \
-      const std::optional<array>& biases,      \
-      const std::optional<array>& lhs_indices, \
-      const std::optional<array>& rhs_indices, \
-      array& out,                              \
-      int bits,                                \
-      int group_size,                          \
-      QuantizationMode mode,                   \
-      cu::CommandEncoder& encoder);            \
-  }
+} // namespace mlx::core

mlx/backend/cuda/quantized/qmm/qmm_sm80.cu

@@ -48,7 +48,9 @@ __global__ void qmm_sm80_kernel(
   CUTE_STATIC_ASSERT_V(congruent(select<0,1,3>(shape_MNKL), dC));
 
   int thread_idx = int(threadIdx.x);
-  auto [m_coord, n_coord, l_coord] = static_cast<uint3>(blockIdx);
+  int m_coord = int(blockIdx.x);
+  int n_coord = int(blockIdx.y);
+  int l_coord = int(blockIdx.z);
 
   // For gather, use index lookup for input batch slicing.
   uint32_t a_batch = lhs_indices ? lhs_indices[l_coord] : l_coord;
@@ -434,7 +436,7 @@ inline void dispatch_quant_types(
 }
 
 template <int TileM>
-void qmm_impl_sm80(
+void qmm_sm80_impl(
     const array& x,
     const array& w,
     const array& scales,
@@ -499,20 +501,19 @@ void qmm_impl_sm80(
   });
 }
 
-} // namespace mlx::core
+// clang-format off
+template void qmm_sm80_impl<@TileM@>(
+    const array& x,
+    const array& w,
+    const array& scales,
+    const std::optional<array>& biases,
+    const std::optional<array>& lhs_indices,
+    const std::optional<array>& rhs_indices,
+    array& out,
+    int bits,
+    int group_size,
+    QuantizationMode mode,
+    cu::CommandEncoder& encoder);
+// clang-format on
 
-#define QMM_SM80_GPU(TileM)                    \
-  namespace mlx::core {                        \
-  template void qmm_impl_sm80<TileM>(          \
-      const array& x,                          \
-      const array& w,                          \
-      const array& scales,                     \
-      const std::optional<array>& biases,      \
-      const std::optional<array>& lhs_indices, \
-      const std::optional<array>& rhs_indices, \
-      array& out,                              \
-      int bits,                                \
-      int group_size,                          \
-      QuantizationMode mode,                   \
-      cu::CommandEncoder& encoder);            \
-  }
+} // namespace mlx::core

mlx/backend/cuda/quantized/qmm/qmm_sm90.cu

@@ -20,8 +20,7 @@ namespace cutlass_gemm {
 using namespace cute;
 
 template <
-    typename TileShapeMN = Shape<_128, _16>,
-    typename ClusterShape = Shape<_1, _1, _1>,
+    int TileN = 16,
     typename Element,
     typename Quant,
     typename GroupSize,
@@ -47,7 +46,8 @@ void qmm_sm90(
 
   using Arch = cutlass::arch::Sm90;
   using Accumulator = float;
-  using TileShape = decltype(append(TileShapeMN{}, Int<kTileShapeK>{}));
+  using TileShape = Shape<_128, Int<TileN>, Int<kTileShapeK>>;
+  using ClusterShape = Shape<Int<(TileN <= 32) ? 1 : 2>, _1, _1>;
 
   using Epilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
       Arch,
@@ -177,8 +177,8 @@ inline void dispatch_groups(int group_size, const char* tag, F&& f) {
   }
 }
 
-template <typename TileShapeMN, typename ClusterShape>
-void qmm_impl_sm90(
+template <int TileN>
+void qmm_sm90_impl(
     const array& x,
     const array& w,
     const array& scales_,
@@ -207,7 +207,7 @@ void qmm_impl_sm90(
         encoder.set_input_array(scales);
         encoder.set_input_array(biases);
         encoder.set_output_array(out);
-        cutlass_gemm::qmm_sm90(
+        cutlass_gemm::qmm_sm90<TileN>(
             gpu_ptr<Element>(x),
             gpu_ptr<Quant>(w),
             gpu_ptr<Element>(scales),
@@ -238,24 +238,19 @@ void qmm_impl_sm90(
   });
 }
 
+// clang-format off
+template void qmm_sm90_impl<@TileN@>(
+    const array& x,
+    const array& w,
+    const array& scales,
+    const array& biases,
+    array& out,
+    int bits,
+    int group_size,
+    cu::CommandEncoder& encoder,
+    Stream s);
+// clang-format on
+
 } // namespace mlx::core
 
-#define QMM_SM90_GPU(TileShapeMN, ClusterShape)           \
-  namespace mlx::core {                                   \
-  template void qmm_impl_sm90<TileShapeMN, ClusterShape>( \
-      const array& x,                                     \
-      const array& w,                                     \
-      const array& scales,                                \
-      const array& biases,                                \
-      array& out,                                         \
-      int bits,                                           \
-      int group_size,                                     \
-      cu::CommandEncoder& encoder,                        \
-      Stream s);                                          \
-  }
-
-#else
-
-#define QMM_SM90_GPU(TileShapeMN, ClusterShape)
-
 #endif // defined(MLX_CUDA_SM90A_ENABLED)

mlx/backend/cuda/quantized/quantized.cpp

@@ -17,7 +17,7 @@ void QuantizedMatmul::eval_gpu(const std::vector<array>& inputs, array& out) {
   auto& s = stream();
   auto& encoder = cu::get_command_encoder(s);
 
-  const array& x = inputs[0];
+  array x = ensure_row_contiguous(inputs[0], encoder, s);
   const array& w = inputs[1];
   const array& scales = inputs[2];
   std::optional<array> biases;
@@ -146,15 +146,17 @@ void GatherQMM::eval_gpu(const std::vector<array>& inputs, array& out) {
   auto& s = stream();
   auto& encoder = cu::get_command_encoder(s);
 
-  const array& x = inputs[0];
+  array x = ensure_row_contiguous(inputs[0], encoder, s);
   const array& w = inputs[1];
   const array& scales = inputs[2];
   std::optional<array> biases;
   if (inputs.size() == 6) {
     biases = inputs[3];
   }
-  array lhs_indices = ensure_contiguous(inputs[inputs.size() - 2], encoder, s);
-  array rhs_indices = ensure_contiguous(inputs[inputs.size() - 1], encoder, s);
+  array lhs_indices =
+      ensure_row_contiguous(inputs[inputs.size() - 2], encoder, s);
+  array rhs_indices =
+      ensure_row_contiguous(inputs[inputs.size() - 1], encoder, s);
 
   int M = out.ndim() > 1 ? out.shape(-2) : 1;
   int N = out.shape(-1);

mlx/distributed/jaccl/lib/README.md

@@ -29,6 +29,8 @@ in macOS 26.2.
 - **Point-to-Point Operations**:
   - `send`: Send data to a specific node
   - `recv`: Receive data from a specific node
+- **Synchronization**:
+  - `barrier`: Block until all nodes in the group reach this point
 - **Type Support**: Bool, Int8-64, UInt8-64, Float16, BFloat16, Float32,
   Float64, Complex64
 
@@ -286,6 +288,9 @@ class Group {
   // Simple send/recv primitives.
   virtual void send(const void* input, size_t n_bytes, int dst) = 0;
   virtual void recv(void* output, size_t n_bytes, int src) = 0;
+
+  // Block until every rank reaches this point.
+  virtual void barrier() = 0;
 };
 ```
 

mlx/distributed/jaccl/lib/examples/CMakeLists.txt

@@ -35,6 +35,7 @@ endfunction()
 # Examples
 build_example(minimal_env.cpp)
 build_example(minimal_cfg.cpp)
+build_example(minimal_barrier.cpp)
 
 # Benchmarks
 build_example(allreduce_bench.cpp)

mlx/distributed/jaccl/lib/examples/minimal_barrier.cpp

@@ -0,0 +1,42 @@
+// Copyright © 2026 Apple Inc.
+//
+// Exercises Group::barrier(). Ranks arrive at the barrier at staggered times;
+// after the barrier returns we do a small all_sum to confirm the group is
+// healthy and that barrier() carried the correct fence semantics.
+
+#include <chrono>
+#include <iostream>
+#include <thread>
+
+#include <jaccl/jaccl.h>
+
+int main() {
+  auto group = jaccl::init();
+  if (!group) {
+    std::cerr << "Failed to initialize JACCL" << std::endl;
+    return 1;
+  }
+
+  int rank = group->rank();
+  int size = group->size();
+
+  std::this_thread::sleep_for(std::chrono::milliseconds(100 * rank));
+  std::cout << "rank " << rank << " entering barrier" << std::endl;
+
+  group->barrier();
+
+  std::cout << "rank " << rank << " exited barrier" << std::endl;
+
+  int in = rank + 1;
+  int out = 0;
+  group->all_sum(&in, &out, sizeof(in), jaccl::Int32);
+  int expected = size * (size + 1) / 2;
+  if (out != expected) {
+    std::cerr << "rank " << rank << ": post-barrier all_sum mismatch (got "
+              << out << ", expected " << expected << ")" << std::endl;
+    return 1;
+  }
+  std::cout << "rank " << rank << ": post-barrier all_sum OK (" << out << ")"
+            << std::endl;
+  return 0;
+}

mlx/distributed/jaccl/lib/jaccl/group.h

@@ -30,6 +30,7 @@ class Group {
 
   virtual void send(const void* input, size_t n_bytes, int dst) = 0;
   virtual void recv(void* output, size_t n_bytes, int src) = 0;
+  virtual void barrier() = 0;
 };
 
 /**

mlx/distributed/jaccl/lib/jaccl/mesh.cpp

@@ -184,6 +184,11 @@ void MeshGroup::recv(void* output, size_t n_bytes, int src) {
   mesh_.recv(static_cast<char*>(output), n_bytes, src);
 }
 
+void MeshGroup::barrier() {
+  uint8_t b = 0;
+  all_sum(&b, &b, sizeof(b), Dtype::UInt8);
+}
+
 template <typename T, typename ReduceOp>
 void MeshGroup::all_reduce(
     const void* input,

mlx/distributed/jaccl/lib/jaccl/mesh.h

@@ -47,6 +47,8 @@ class MeshGroup : public Group {
   void send(const void* input, size_t n_bytes, int dst) override;
   void recv(void* output, size_t n_bytes, int src) override;
 
+  void barrier() override;
+
  private:
   template <typename T, typename ReduceOp>
   void all_reduce(

mlx/distributed/jaccl/lib/jaccl/ring.cpp

@@ -190,6 +190,11 @@ void RingGroup::recv(void* output, size_t n_bytes, int src) {
   ring_.recv(static_cast<char*>(output), n_bytes, src, n_conns_);
 }
 
+void RingGroup::barrier() {
+  uint8_t b = 0;
+  all_sum(&b, &b, sizeof(b), Dtype::UInt8);
+}
+
 template <typename T, typename ReduceOp>
 void RingGroup::all_reduce(
     const void* input,

mlx/distributed/jaccl/lib/jaccl/ring.h

@@ -48,6 +48,8 @@ class RingGroup : public Group {
   void send(const void* input, size_t n_bytes, int dst) override;
   void recv(void* output, size_t n_bytes, int src) override;
 
+  void barrier() override;
+
  private:
   template <typename T, typename ReduceOp>
   void all_reduce(

mlx/io/CMakeLists.txt

@@ -17,6 +17,11 @@ if(MLX_BUILD_GGUF)
                              PRIVATE $<BUILD_INTERFACE:${gguflib_SOURCE_DIR}>)
   add_library(gguflib STATIC ${gguflib_SOURCE_DIR}/fp16.c
                              ${gguflib_SOURCE_DIR}/gguflib.c)
+  # gguflib uses assert() to reject malformed tensor headers (e.g. ndim > 8).
+  # Those checks are otherwise compiled out by -DNDEBUG in release builds, which
+  # leaves out-of-bounds reads/writes unguarded when loading untrusted GGUF
+  # files. Force NDEBUG off for this target so the asserts stay live.
+  target_compile_options(gguflib PRIVATE -UNDEBUG)
   target_link_libraries(mlx PRIVATE $<BUILD_INTERFACE:gguflib>)
   target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/gguf.cpp
                              ${CMAKE_CURRENT_SOURCE_DIR}/gguf_quants.cpp)

mlx/io/safetensors.cpp

@@ -28,6 +28,7 @@ using json = nlohmann::json;
 #define ST_U32 "U32"
 #define ST_U64 "U64"
 #define ST_F8_E4M3 "F8_E4M3"
+#define ST_F8_E8M0 "F8_E8M0"
 
 // Note: Complex numbers aren't in the spec yet so this could change -
 // https://github.com/huggingface/safetensors/issues/389
@@ -97,6 +98,8 @@ Dtype dtype_from_safetensor_str(std::string_view str) {
     return complex64;
   } else if (str == ST_F8_E4M3) {
     return uint8;
+  } else if (str == ST_F8_E8M0) {
+    return uint8;
   } else {
     std::ostringstream msg;
     msg << "[safetensor] unsupported dtype" << str;

mlx/linalg.cpp

@@ -705,4 +705,170 @@ array solve_triangular(
   return matmul(a_inv, b, s);
 }
 
-} // namespace mlx::core::linalg
+void validate_det(
+    const array& a,
+    const StreamOrDevice& stream,
+    const std::string& fname) {
+  check_cpu_stream(stream, fname);
+  if (issubdtype(a.dtype(), complexfloating)) {
+    throw std::invalid_argument(fname + " Complex inputs are not supported.");
+  }
+  if (a.ndim() < 2) {
+    std::ostringstream msg;
+    msg << fname
+        << " Arrays must have >= 2 dimensions. Received array "
+           "with "
+        << a.ndim() << " dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+
+  if (a.shape(-1) != a.shape(-2)) {
+    throw std::invalid_argument(fname + " Only defined for square matrices.");
+  }
+}
+
+array det_raw_small(const array& a, StreamOrDevice s) {
+  int n = a.shape(-1);
+
+  // Empty 0x0 matrix: determinant is the empty product = 1
+  if (n == 0) {
+    Shape out_shape(a.shape().begin(), a.shape().end() - 2);
+    return broadcast_to(array(1.0f, a.dtype()), std::move(out_shape), s);
+  }
+
+  // Helper to extract a[..., i, j] from the last two dims
+  auto elem = [&](int i, int j) {
+    auto starts = Shape(a.ndim(), 0);
+    auto stops = a.shape();
+    starts[a.ndim() - 2] = i;
+    stops[a.ndim() - 2] = i + 1;
+    starts[a.ndim() - 1] = j;
+    stops[a.ndim() - 1] = j + 1;
+    return squeeze(squeeze(slice(a, starts, stops, s), -1, s), -1, s);
+  };
+
+  if (n == 1) {
+    return elem(0, 0);
+  } else if (n == 2) {
+    return subtract(
+        multiply(elem(0, 0), elem(1, 1), s),
+        multiply(elem(0, 1), elem(1, 0), s),
+        s);
+  } else {
+    // 3x3: a00*(a11*a22 - a12*a21) - a01*(a10*a22 - a12*a20) + a02*(a10*a21 -
+    // a11*a20)
+    auto a00 = elem(0, 0), a01 = elem(0, 1), a02 = elem(0, 2);
+    auto a10 = elem(1, 0), a11 = elem(1, 1), a12 = elem(1, 2);
+    auto a20 = elem(2, 0), a21 = elem(2, 1), a22 = elem(2, 2);
+    return add(
+        subtract(
+            multiply(
+                a00,
+                subtract(multiply(a11, a22, s), multiply(a12, a21, s), s),
+                s),
+            multiply(
+                a01,
+                subtract(multiply(a10, a22, s), multiply(a12, a20, s), s),
+                s),
+            s),
+        multiply(
+            a02, subtract(multiply(a10, a21, s), multiply(a11, a20, s), s), s),
+        s);
+  }
+}
+
+std::pair<array, array> slogdet_impl(const array& input, StreamOrDevice s) {
+  int n = input.shape(-1);
+  auto dtype = input.dtype();
+
+  // Small-matrix fast path
+  if (n <= 3) {
+    auto raw = det_raw_small(input, s);
+    auto abs_raw = abs(raw, s);
+    auto sgn = sign(raw, s);
+    auto logabs = log(abs_raw, s);
+    return std::make_pair(sgn, logabs);
+  }
+
+  // General LU-based path
+  auto [LU, pivots] = lu_factor(input, s);
+
+  // Extract diagonal of U
+  auto diag = diagonal(LU, 0, -2, -1, s);
+
+  // Permutation parity: count positions where pivot[i] != i
+  int k = std::min(input.shape(-2), input.shape(-1));
+  auto iota = arange(0, k, uint32, s);
+  auto parity = astype(
+      sum(not_equal(pivots, iota, s),
+          /* axis = */ -1,
+          /* keepdims = */ false,
+          s),
+      int32,
+      s);
+
+  // Count negative diagonal elements
+  auto num_neg = astype(
+      sum(less(diag, array(0.0f, dtype), s),
+          /* axis = */ -1,
+          /* keepdims = */ false,
+          s),
+      int32,
+      s);
+
+  // sign = (-1)^(parity + num_neg)
+  auto total = add(parity, num_neg, s);
+  auto sign_val = astype(
+      subtract(
+          array(1, int32),
+          multiply(array(2, int32), remainder(total, array(2, int32), s), s),
+          s),
+      dtype,
+      s);
+
+  // logabsdet = sum(log(abs(diag)))
+  auto logabsdet =
+      sum(log(abs(diag, s), s), /* axis = */ -1, /* keepdims = */ false, s);
+
+  // Handle singular matrices: any zero on diagonal
+  auto is_zero =
+      any(equal(diag, array(0.0f, dtype), s),
+          /* axis = */ -1,
+          /* keepdims = */ false,
+          s);
+  sign_val = where(is_zero, array(0.0f, dtype), sign_val, s);
+  logabsdet = where(
+      is_zero,
+      array(-std::numeric_limits<float>::infinity(), dtype),
+      logabsdet,
+      s);
+
+  return std::make_pair(sign_val, logabsdet);
+}
+
+std::pair<array, array> slogdet(const array& a, StreamOrDevice s /* = {} */) {
+  validate_det(a, s, "[linalg::slogdet]");
+
+  auto dtype = at_least_float(a.dtype());
+  auto input = astype(a, dtype, s);
+  return slogdet_impl(input, s);
+}
+
+array det(const array& a, StreamOrDevice s /* = {} */) {
+  validate_det(a, s, "[linalg::det]");
+
+  auto dtype = at_least_float(a.dtype());
+  auto input = astype(a, dtype, s);
+  int n = input.shape(-1);
+
+  // Small-matrix fast path: compute directly, skip log/exp round-trip
+  if (n <= 3) {
+    return det_raw_small(input, s);
+  }
+
+  // General case: det = sign * exp(logabsdet)
+  auto [sign_val, logabsdet] = slogdet_impl(input, s);
+  return multiply(sign_val, exp(logabsdet, s), s);
+}
+
+} // namespace mlx::core::linalg

mlx/linalg.h

@@ -112,4 +112,8 @@ eigvalsh(const array& a, std::string UPLO = "L", StreamOrDevice s = {});
 MLX_API std::pair<array, array>
 eigh(const array& a, std::string UPLO = "L", StreamOrDevice s = {});
 
+MLX_API array det(const array& a, StreamOrDevice s = {});
+
+MLX_API std::pair<array, array> slogdet(const array& a, StreamOrDevice s = {});
+
 } // namespace mlx::core::linalg

mlx/version.h

@@ -5,8 +5,8 @@
 #include "mlx/api.h"
 
 #define MLX_VERSION_MAJOR 0
-#define MLX_VERSION_MINOR 31
-#define MLX_VERSION_PATCH 2
+#define MLX_VERSION_MINOR 32
+#define MLX_VERSION_PATCH 0
 #define MLX_VERSION_NUMERIC \
   (100000 * MLX_VERSION_MAJOR + 1000 * MLX_VERSION_MINOR + MLX_VERSION_PATCH)
 

python/src/CMakeLists.txt

@@ -2,6 +2,7 @@ nanobind_add_module(
   core
   NB_STATIC
   STABLE_ABI
+  FREE_THREADED
   LTO
   NOMINSIZE
   NB_DOMAIN

python/src/convert.cpp

@@ -1,5 +1,8 @@
 // Copyright © 2024 Apple Inc.
 
+#include <limits>
+#include <sstream>
+
 #include <nanobind/stl/complex.h>
 
 #include "python/src/convert.h"
@@ -15,9 +18,15 @@ enum PyScalarT {
 };
 
 int check_shape_dim(int64_t dim) {
-  if (dim > std::numeric_limits<int>::max()) {
-    throw std::invalid_argument(
-        "Shape dimension falls outside supported `int` range.");
+  if (dim > std::numeric_limits<int>::max() ||
+      dim < std::numeric_limits<int>::min()) {
+    std::ostringstream msg;
+    msg << "Shape dimension " << dim << " is outside the supported range ["
+        << std::numeric_limits<int>::min() << ", "
+        << std::numeric_limits<int>::max()
+        << "]. MLX currently uses 32-bit integers for shape dimensions.";
+    PyErr_SetString(PyExc_OverflowError, msg.str().c_str());
+    nb::detail::raise_python_error();
   }
   return static_cast<int>(dim);
 }

python/src/convert.h

@@ -76,3 +76,7 @@ nb::object tolist(mx::array& a);
 mx::array create_array(nb::object v, std::optional<mx::Dtype> t);
 mx::array array_from_list(nb::list pl, std::optional<mx::Dtype> dtype);
 mx::array array_from_list(nb::tuple pl, std::optional<mx::Dtype> dtype);
+
+// Narrow a Python-side shape dimension (int64) to a C++ mx::ShapeElem (int32),
+// raising a clear error if the value would overflow.
+int check_shape_dim(int64_t dim);

python/src/linalg.cpp

@@ -660,4 +660,77 @@ void init_linalg(nb::module_& parent_module) {
         Returns:
             array: The unique solution to the system ``AX = B``.
       )pbdoc");
+
+  m.def(
+      "det",
+      &mx::linalg::det,
+      "a"_a,
+      nb::kw_only(),
+      "stream"_a = nb::none(),
+      nb::sig(
+          "def det(a: array, *, stream: Union[None, Stream, Device] = None) -> array"),
+      R"pbdoc(
+        Compute the determinant of a square matrix.
+
+        This function supports arrays with at least 2 dimensions. When the
+        input has more than two dimensions, the determinant is computed for
+        each matrix in the last two dimensions.
+
+        Args:
+            a (array): Input array.
+            stream (Stream, optional): Stream or device. Defaults to ``None``
+              in which case the default stream of the default device is used.
+
+        Returns:
+            array: The determinant(s) of the input matrix (matrices).
+
+        Example:
+            >>> A = mx.array([[1., 2.], [3., 4.]])
+            >>> mx.linalg.det(A, stream=mx.cpu)
+            array(-2, dtype=float32)
+      )pbdoc");
+
+  m.def(
+      "slogdet",
+      [](const mx::array& a, mx::StreamOrDevice s) {
+        auto result = mx::linalg::slogdet(a, s);
+        return nb::make_tuple(result.first, result.second);
+      },
+      "a"_a,
+      nb::kw_only(),
+      "stream"_a = nb::none(),
+      nb::sig(
+          "def slogdet(a: array, *, stream: Union[None, Stream, Device] = None) -> Tuple[array, array]"),
+      R"pbdoc(
+        Compute the sign and natural log of the absolute value of the
+        determinant of a square matrix.
+
+        This function supports arrays with at least 2 dimensions. When the
+        input has more than two dimensions, the sign and log-absolute-determinant
+        are computed for each matrix in the last two dimensions.
+
+        For a singular matrix, ``sign`` is 0 and ``logabsdet`` is ``-inf``.
+
+        The determinant can be reconstructed as ``det = sign * exp(logabsdet)``.
+        This is more numerically stable than computing the determinant directly
+        for matrices with large or small determinants.
+
+        Args:
+            a (array): Input array.
+            stream (Stream, optional): Stream or device. Defaults to ``None``
+              in which case the default stream of the default device is used.
+
+        Returns:
+            tuple(array, array): The ``sign`` and ``logabsdet`` of the
+              determinant. ``sign`` is -1, 0, or +1. ``logabsdet`` is the
+              natural log of the absolute value of the determinant.
+
+        Example:
+            >>> A = mx.array([[1., 2.], [3., 4.]])
+            >>> sign, logabsdet = mx.linalg.slogdet(A, stream=mx.cpu)
+            >>> sign
+            array(-1, dtype=float32)
+            >>> logabsdet
+            array(0.693147, dtype=float32)
+      )pbdoc");
 }

python/src/ops.cpp

@@ -15,6 +15,7 @@
 #include "mlx/einsum.h"
 #include "mlx/ops.h"
 #include "mlx/utils.h"
+#include "python/src/convert.h"
 #include "python/src/load.h"
 #include "python/src/small_vector.h"
 #include "python/src/utils.h"
@@ -45,6 +46,13 @@ double scalar_to_double(Scalar s) {
   }
 }
 
+mx::Shape to_shape(const nb::object& shape) {
+  if (nb::isinstance<nb::int_>(shape)) {
+    return {check_shape_dim(nb::cast<int64_t>(shape))};
+  }
+  return nb::cast<mx::Shape>(shape);
+}
+
 void init_ops(nb::module_& m) {
   m.def(
       "reshape",
@@ -1702,15 +1710,11 @@ void init_ops(nb::module_& m) {
       )pbdoc");
   m.def(
       "full",
-      [](const std::variant<int, mx::Shape>& shape,
+      [](const nb::object& shape,
          const ScalarOrArray& vals,
          std::optional<mx::Dtype> dtype,
          mx::StreamOrDevice s) {
-        if (auto pv = std::get_if<int>(&shape); pv) {
-          return mx::full({*pv}, to_array(vals, dtype), s);
-        } else {
-          return mx::full(std::get<mx::Shape>(shape), to_array(vals, dtype), s);
-        }
+        return mx::full(to_shape(shape), to_array(vals, dtype), s);
       },
       "shape"_a,
       "vals"_a,
@@ -1736,15 +1740,11 @@ void init_ops(nb::module_& m) {
       )pbdoc");
   m.def(
       "zeros",
-      [](const std::variant<int, mx::Shape>& shape,
+      [](const nb::object& shape,
          std::optional<mx::Dtype> dtype,
          mx::StreamOrDevice s) {
         auto t = dtype.value_or(mx::float32);
-        if (auto pv = std::get_if<int>(&shape); pv) {
-          return mx::zeros({*pv}, t, s);
-        } else {
-          return mx::zeros(std::get<mx::Shape>(shape), t, s);
-        }
+        return mx::zeros(to_shape(shape), t, s);
       },
       "shape"_a,
       "dtype"_a.none() = mx::float32,
@@ -1802,15 +1802,11 @@ void init_ops(nb::module_& m) {
       )pbdoc");
   m.def(
       "ones",
-      [](const std::variant<int, mx::Shape>& shape,
+      [](const nb::object& shape,
          std::optional<mx::Dtype> dtype,
          mx::StreamOrDevice s) {
         auto t = dtype.value_or(mx::float32);
-        if (auto pv = std::get_if<int>(&shape); pv) {
-          return mx::ones({*pv}, t, s);
-        } else {
-          return mx::ones(std::get<mx::Shape>(shape), t, s);
-        }
+        return mx::ones(to_shape(shape), t, s);
       },
       "shape"_a,
       "dtype"_a.none() = mx::float32,

python/src/small_vector.h

@@ -2,6 +2,11 @@
 
 #pragma once
 
+#include <cstdint>
+#include <limits>
+#include <sstream>
+#include <type_traits>
+
 #include "mlx/small_vector.h"
 
 #include <nanobind/stl/detail/nb_list.h>
@@ -14,11 +19,19 @@ struct type_caster<mlx::core::SmallVector<Type, Size, Alloc>> {
   using List = mlx::core::SmallVector<Type, Size, Alloc>;
   using Caster = make_caster<Type>;
 
+  // For narrow integer element types we fetch each element through a wider
+  // integer caster so we can emit a clean OverflowError on overflow instead of
+  // nanobind's generic "incompatible function arguments" TypeError.
+  static constexpr bool kNarrowInt = std::is_integral_v<Type> &&
+      !std::is_same_v<Type, bool> && (sizeof(Type) < sizeof(int64_t));
+
   NB_TYPE_CASTER(
       List,
       const_name("tuple[") + make_caster<Type>::Name + const_name(", ...]"))
 
-  bool from_python(handle src, uint8_t flags, cleanup_list* cleanup) noexcept {
+  // Not noexcept: on overflow of a narrow integer element we raise
+  // OverflowError so nanobind surfaces a clean error to the user.
+  bool from_python(handle src, uint8_t flags, cleanup_list* cleanup) {
     size_t size;
     PyObject* temp;
 
@@ -29,19 +42,39 @@ struct type_caster<mlx::core::SmallVector<Type, Size, Alloc>> {
     value.clear();
     value.reserve(size);
 
-    Caster caster;
     bool success = o != nullptr;
 
     flags = flags_for_local_caster<Type>(flags);
 
     for (size_t i = 0; i < size; ++i) {
-      if (!caster.from_python(o[i], flags, cleanup) ||
-          !caster.template can_cast<Type>()) {
-        success = false;
-        break;
+      if constexpr (kNarrowInt) {
+        make_caster<int64_t> wide;
+        if (!wide.from_python(o[i], flags, cleanup) ||
+            !wide.template can_cast<int64_t>()) {
+          success = false;
+          break;
+        }
+        int64_t v = wide.operator cast_t<int64_t>();
+        if (v > std::numeric_limits<Type>::max() ||
+            v < std::numeric_limits<Type>::min()) {
+          std::ostringstream msg;
+          msg << "Integer value " << v << " is outside the supported range ["
+              << static_cast<int64_t>(std::numeric_limits<Type>::min()) << ", "
+              << static_cast<int64_t>(std::numeric_limits<Type>::max()) << "].";
+          Py_XDECREF(temp);
+          PyErr_SetString(PyExc_OverflowError, msg.str().c_str());
+          raise_python_error();
+        }
+        value.push_back(static_cast<Type>(v));
+      } else {
+        Caster caster;
+        if (!caster.from_python(o[i], flags, cleanup) ||
+            !caster.template can_cast<Type>()) {
+          success = false;
+          break;
+        }
+        value.push_back(caster.operator cast_t<Type>());
       }
-
-      value.push_back(caster.operator cast_t<Type>());
     }
 
     Py_XDECREF(temp);

python/tests/cuda_skip.py

@@ -1,22 +1,5 @@
 cuda_skip = {
-    # Lapack ops NYI
-    "TestLinalg.test_cholesky",
-    "TestLinalg.test_cholesky_inv",
-    "TestLinalg.test_eig",
-    "TestLinalg.test_eigh",
-    "TestLinalg.test_inverse",
-    "TestVmap.test_vmap_inverse",
-    "TestLinalg.test_lu",
-    "TestLinalg.test_lu_factor",
-    "TestLinalg.test_pseudo_inverse",
-    "TestLinalg.test_qr_factorization",
-    "TestInit.test_orthogonal",
-    "TestLinalg.test_svd_decomposition",
-    "TestVmap.test_vmap_svd",
-    "TestLinalg.test_tri_inverse",
     # Quantization NYI
     "TestQuantized.test_gather_matmul_grad",
     "TestQuantized.test_gather_qmm",
-    "TestQuantized.test_gather_qmm_sorted",
-    "TestQuantized.test_gather_qmm_grad",
 }

python/tests/test_array.py

@@ -767,10 +767,13 @@ class TestArray(mlx_tests.MLXTestCase):
 
     def test_array_np_shape_dim_check(self):
         a_npy = np.empty(2**31, dtype=np.bool_)
-        with self.assertRaises(ValueError) as e:
+        with self.assertRaises(OverflowError) as e:
             mx.array(a_npy)
         self.assertEqual(
-            str(e.exception), "Shape dimension falls outside supported `int` range."
+            str(e.exception),
+            "Shape dimension 2147483648 is outside the supported range "
+            "[-2147483648, 2147483647]. MLX currently uses 32-bit integers "
+            "for shape dimensions.",
         )
 
     def test_dtype_promotion(self):

python/tests/test_linalg.py

@@ -520,6 +520,19 @@ class TestLinalg(mlx_tests.MLXTestCase):
         P, L, U = mx.linalg.lu(a, stream=mx.cpu)
         self.assertTrue(mx.allclose(L[P, :] @ U, a))
 
+        # Test singular matrix (should not throw)
+        a = mx.array(
+            [
+                [1.0, 2.0, 3.0, 4.0],
+                [2.0, 4.0, 6.0, 8.0],
+                [0.0, 1.0, 1.0, 0.0],
+                [1.0, 0.0, 0.0, 1.0],
+            ]
+        )
+        P, L, U = mx.linalg.lu(a, stream=mx.cpu)
+        L_permuted = mx.take_along_axis(L, P[..., None], axis=-2)
+        self.assertTrue(mx.allclose(L_permuted @ U, a))
+
     def test_lu_factor(self):
         mx.random.seed(7)
 
@@ -616,6 +629,248 @@ class TestLinalg(mlx_tests.MLXTestCase):
         expected = np.linalg.solve(a, b)
         self.assertTrue(np.allclose(result, expected))
 
+    def test_det(self):
+        # 1x1 fast path
+        A = mx.array([[5.0]])
+        self.assertTrue(np.allclose(mx.linalg.det(A, stream=mx.cpu), 5.0))
+
+        # 2x2 fast path
+        A = mx.array([[1.0, 2.0], [3.0, 4.0]])
+        d = mx.linalg.det(A, stream=mx.cpu)
+        self.assertTrue(np.allclose(d, -2.0))
+
+        # 3x3 fast path
+        A = mx.array([[1.0, 2.0, 3.0], [0.0, 1.0, 4.0], [5.0, 6.0, 0.0]])
+        d = mx.linalg.det(A, stream=mx.cpu)
+        expected = np.linalg.det(np.array(A))
+        self.assertTrue(np.allclose(d, expected, atol=1e-5))
+
+        # 4x4 LU path: compare with numpy
+        np.random.seed(42)
+        A_np = np.random.randn(4, 4).astype(np.float32)
+        A_mx = mx.array(A_np)
+        d_mx = mx.linalg.det(A_mx, stream=mx.cpu)
+        d_np = np.linalg.det(A_np)
+        self.assertTrue(np.allclose(d_mx, d_np, atol=1e-4))
+
+        # 5x5 LU path
+        A_np = np.random.randn(5, 5).astype(np.float32)
+        A_mx = mx.array(A_np)
+        d_mx = mx.linalg.det(A_mx, stream=mx.cpu)
+        d_np = np.linalg.det(A_np)
+        self.assertTrue(np.allclose(d_mx, d_np, atol=1e-4))
+
+        # Identity matrix
+        A = mx.eye(5)
+        self.assertTrue(np.allclose(mx.linalg.det(A, stream=mx.cpu), 1.0))
+
+        # Batched: (3, 4, 4)
+        A_np = np.random.randn(3, 4, 4).astype(np.float32)
+        A_mx = mx.array(A_np)
+        d_mx = mx.linalg.det(A_mx, stream=mx.cpu)
+        d_np = np.linalg.det(A_np)
+        self.assertTrue(np.allclose(d_mx, d_np, atol=1e-4))
+
+        # Multi-batch: (2, 3, 3, 3)
+        A_np = np.random.randn(2, 3, 3, 3).astype(np.float32)
+        A_mx = mx.array(A_np)
+        d_mx = mx.linalg.det(A_mx, stream=mx.cpu)
+        d_np = np.linalg.det(A_np)
+        self.assertTrue(np.allclose(d_mx, d_np, atol=1e-4))
+
+        # Integer input auto-promotes to float
+        A = mx.array([[1, 2], [3, 4]])
+        d = mx.linalg.det(A, stream=mx.cpu)
+        self.assertTrue(np.allclose(d, -2.0))
+
+        # float64
+        A_np = np.random.randn(4, 4).astype(np.float64)
+        A_mx = mx.array(A_np)
+        d_mx = mx.linalg.det(A_mx, stream=mx.cpu)
+        d_np = np.linalg.det(A_np)
+        self.assertTrue(np.allclose(d_mx, d_np, atol=1e-10))
+
+        # Singular 4x4 matrix (LU path): det should be 0
+        A = mx.array(
+            [
+                [1.0, 2.0, 3.0, 4.0],
+                [2.0, 4.0, 6.0, 8.0],
+                [0.0, 1.0, 1.0, 0.0],
+                [1.0, 0.0, 0.0, 1.0],
+            ]
+        )
+        d = mx.linalg.det(A, stream=mx.cpu)
+        self.assertTrue(np.allclose(d, 0.0, atol=1e-5))
+
+        # Singular 5x5 matrix (LU path)
+        A_np = np.ones((5, 5), dtype=np.float32)
+        A_mx = mx.array(A_np)
+        d = mx.linalg.det(A_mx, stream=mx.cpu)
+        self.assertTrue(np.allclose(d, 0.0, atol=1e-5))
+
+        # Batched singular matrices (LU path)
+        A_np = np.array([np.diag([1.0, 2.0, 0.0, 3.0]), np.eye(4, dtype=np.float32)])
+        A_mx = mx.array(A_np)
+        d_mx = mx.linalg.det(A_mx, stream=mx.cpu)
+        d_np = np.linalg.det(A_np)
+        self.assertTrue(np.allclose(d_mx, d_np, atol=1e-5))
+
+        # Empty 0x0 matrix: det is the empty product = 1
+        d = mx.linalg.det(mx.zeros((0, 0)), stream=mx.cpu)
+        self.assertEqual(d.shape, ())
+        self.assertEqual(float(d), 1.0)
+
+        # Batched empty matrices: shape preserves batch dims
+        d = mx.linalg.det(mx.zeros((3, 0, 0)), stream=mx.cpu)
+        self.assertTrue(np.allclose(d, np.linalg.det(np.zeros((3, 0, 0)))))
+
+        # Error: non-square
+        with self.assertRaises(ValueError):
+            mx.linalg.det(mx.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]), stream=mx.cpu)
+
+        # Error: 1D
+        with self.assertRaises(ValueError):
+            mx.linalg.det(mx.array([1.0, 2.0]), stream=mx.cpu)
+
+        # Error: complex unsupported (small-matrix path)
+        with self.assertRaises(ValueError):
+            mx.linalg.det(mx.array([[1.0 + 1j, 2.0], [3.0, 4.0]]), stream=mx.cpu)
+
+        # Error: complex unsupported (LU path)
+        with self.assertRaises(ValueError):
+            mx.linalg.det(mx.eye(4).astype(mx.complex64), stream=mx.cpu)
+
+    def test_slogdet(self):
+        # 2x2: det = -2 => sign = -1, logabsdet = log(2)
+        A = mx.array([[1.0, 2.0], [3.0, 4.0]])
+        sign, logabsdet = mx.linalg.slogdet(A, stream=mx.cpu)
+        self.assertTrue(np.allclose(sign, -1.0))
+        self.assertTrue(np.allclose(logabsdet, np.log(2.0), atol=1e-5))
+
+        # Identity: sign = 1, logabsdet = 0
+        A = mx.eye(4)
+        sign, logabsdet = mx.linalg.slogdet(A, stream=mx.cpu)
+        self.assertTrue(np.allclose(sign, 1.0))
+        self.assertTrue(np.allclose(logabsdet, 0.0, atol=1e-6))
+
+        # Compare with numpy for random matrices
+        np.random.seed(42)
+        for n in [1, 2, 3, 4, 5]:
+            A_np = np.random.randn(n, n).astype(np.float32)
+            A_mx = mx.array(A_np)
+            sign_mx, logabs_mx = mx.linalg.slogdet(A_mx, stream=mx.cpu)
+            sign_np, logabs_np = np.linalg.slogdet(A_np)
+            with self.subTest(n=n):
+                self.assertTrue(np.allclose(sign_mx, sign_np, atol=1e-5))
+                self.assertTrue(np.allclose(logabs_mx, logabs_np, atol=1e-4))
+
+        # Singular matrix 2x2 (fast path): sign = 0, logabsdet = -inf
+        A = mx.array([[1.0, 2.0], [2.0, 4.0]])
+        sign, logabsdet = mx.linalg.slogdet(A, stream=mx.cpu)
+        self.assertEqual(float(sign), 0.0)
+        self.assertEqual(float(logabsdet), float("-inf"))
+
+        # Singular 4x4 matrix (LU path): sign = 0, logabsdet = -inf
+        A = mx.array(
+            [
+                [1.0, 2.0, 3.0, 4.0],
+                [2.0, 4.0, 6.0, 8.0],
+                [0.0, 1.0, 1.0, 0.0],
+                [1.0, 0.0, 0.0, 1.0],
+            ]
+        )
+        sign, logabsdet = mx.linalg.slogdet(A, stream=mx.cpu)
+        self.assertEqual(float(sign), 0.0)
+        self.assertEqual(float(logabsdet), float("-inf"))
+
+        # Singular 5x5 matrix (LU path): all-ones matrix
+        A = mx.array(np.ones((5, 5), dtype=np.float32))
+        sign, logabsdet = mx.linalg.slogdet(A, stream=mx.cpu)
+        self.assertEqual(float(sign), 0.0)
+        self.assertEqual(float(logabsdet), float("-inf"))
+
+        # Batched with mix of singular and non-singular (LU path)
+        A_np = np.array([np.diag([1.0, 2.0, 0.0, 3.0]), np.eye(4, dtype=np.float32)])
+        A_mx = mx.array(A_np)
+        sign_mx, logabs_mx = mx.linalg.slogdet(A_mx, stream=mx.cpu)
+        sign_np, logabs_np = np.linalg.slogdet(A_np)
+        self.assertTrue(np.allclose(sign_mx, sign_np, atol=1e-5))
+        # Check -inf for singular, 0.0 for identity
+        self.assertEqual(float(logabs_mx[0]), float("-inf"))
+        self.assertTrue(np.allclose(logabs_mx[1], 0.0, atol=1e-6))
+
+        # Batched
+        A_np = np.random.randn(3, 4, 4).astype(np.float32)
+        A_mx = mx.array(A_np)
+        sign_mx, logabs_mx = mx.linalg.slogdet(A_mx, stream=mx.cpu)
+        sign_np, logabs_np = np.linalg.slogdet(A_np)
+        self.assertTrue(np.allclose(sign_mx, sign_np, atol=1e-5))
+        self.assertTrue(np.allclose(logabs_mx, logabs_np, atol=1e-4))
+
+        # Multi-batch
+        A_np = np.random.randn(2, 3, 3, 3).astype(np.float32)
+        A_mx = mx.array(A_np)
+        sign_mx, logabs_mx = mx.linalg.slogdet(A_mx, stream=mx.cpu)
+        sign_np, logabs_np = np.linalg.slogdet(A_np)
+        self.assertTrue(np.allclose(sign_mx, sign_np, atol=1e-5))
+        self.assertTrue(np.allclose(logabs_mx, logabs_np, atol=1e-4))
+
+        # Numerical stability: large matrix where det overflows
+        # 0.1 * I_100 has det = 0.1^100 which underflows in float32
+        # but slogdet should give sign=1, logabsdet = 100*log(0.1)
+        n = 100
+        A = mx.array(0.1) * mx.eye(n)
+        sign, logabsdet = mx.linalg.slogdet(A, stream=mx.cpu)
+        self.assertTrue(np.allclose(sign, 1.0))
+        self.assertTrue(np.allclose(logabsdet, n * np.log(0.1), atol=1e-3))
+
+        # Verify det = sign * exp(logabsdet) for non-singular cases
+        A_np = np.random.randn(5, 5).astype(np.float32)
+        A_mx = mx.array(A_np)
+        sign_mx, logabs_mx = mx.linalg.slogdet(A_mx, stream=mx.cpu)
+        det_mx = mx.linalg.det(A_mx, stream=mx.cpu)
+        reconstructed = float(sign_mx) * np.exp(float(logabs_mx))
+        self.assertTrue(np.allclose(float(det_mx), reconstructed, rtol=1e-4))
+
+        # float64
+        A_np = np.random.randn(4, 4).astype(np.float64)
+        A_mx = mx.array(A_np)
+        sign_mx, logabs_mx = mx.linalg.slogdet(A_mx, stream=mx.cpu)
+        sign_np, logabs_np = np.linalg.slogdet(A_np)
+        self.assertTrue(np.allclose(sign_mx, sign_np))
+        self.assertTrue(np.allclose(logabs_mx, logabs_np, atol=1e-10))
+
+        # Empty 0x0 matrix: sign = 1, logabsdet = 0 (empty product)
+        sign, logabsdet = mx.linalg.slogdet(mx.zeros((0, 0)), stream=mx.cpu)
+        self.assertEqual(sign.shape, ())
+        self.assertEqual(logabsdet.shape, ())
+        self.assertEqual(float(sign), 1.0)
+        self.assertEqual(float(logabsdet), 0.0)
+
+        # Batched empty matrices
+        sign, logabsdet = mx.linalg.slogdet(mx.zeros((3, 0, 0)), stream=mx.cpu)
+        sign_np, logabs_np = np.linalg.slogdet(np.zeros((3, 0, 0)))
+        self.assertTrue(np.allclose(sign, sign_np))
+        self.assertTrue(np.allclose(logabsdet, logabs_np))
+
+        # Error: non-square
+        with self.assertRaises(ValueError):
+            mx.linalg.slogdet(
+                mx.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]), stream=mx.cpu
+            )
+
+        # Error: 1D
+        with self.assertRaises(ValueError):
+            mx.linalg.slogdet(mx.array([1.0, 2.0]), stream=mx.cpu)
+
+        # Error: complex unsupported (small-matrix path)
+        with self.assertRaises(ValueError):
+            mx.linalg.slogdet(mx.array([[1.0 + 1j, 2.0], [3.0, 4.0]]), stream=mx.cpu)
+
+        # Error: complex unsupported (LU path)
+        with self.assertRaises(ValueError):
+            mx.linalg.slogdet(mx.eye(4).astype(mx.complex64), stream=mx.cpu)
+
 
 if __name__ == "__main__":
     mlx_tests.MLXTestRunner()

python/tests/test_ops.py

@@ -92,6 +92,52 @@ class TestOps(mlx_tests.MLXTestCase):
             self.assertEqual(y.dtype, t)
             self.assertTrue(mx.array_equal(y, x))
 
+    def test_shape_overflow_error(self):
+        # Shape dimensions that don't fit in int32 should raise a clear
+        # OverflowError that names the offending value, rather than a generic
+        # "incompatible function arguments" TypeError. The overflow check
+        # lives in the mx::Shape type caster, so it applies to every op that
+        # takes a shape. See issue #2681.
+        too_big = 2**31
+
+        # Array creation ops — also exercise the scalar shape path.
+        for ctor in (mx.zeros, mx.ones):
+            with self.assertRaises(OverflowError) as cm:
+                ctor(too_big)
+            self.assertIn(str(too_big), str(cm.exception))
+            with self.assertRaises(OverflowError) as cm:
+                ctor([too_big])
+            self.assertIn(str(too_big), str(cm.exception))
+
+        with self.assertRaises(OverflowError) as cm:
+            mx.full(too_big, 0.0)
+        self.assertIn(str(too_big), str(cm.exception))
+        with self.assertRaises(OverflowError) as cm:
+            mx.full([too_big], 0.0)
+        self.assertIn(str(too_big), str(cm.exception))
+
+        # Other shape-taking ops should surface the same clean error.
+        a = mx.zeros(4)
+        with self.assertRaises(OverflowError) as cm:
+            mx.reshape(a, [too_big])
+        self.assertIn(str(too_big), str(cm.exception))
+
+        with self.assertRaises(OverflowError) as cm:
+            mx.broadcast_to(a, [too_big, 1])
+        self.assertIn(str(too_big), str(cm.exception))
+
+        # Negative overflow (< int32 min) is caught too.
+        too_negative = -(2**31) - 1
+        with self.assertRaises(OverflowError) as cm:
+            mx.zeros([too_negative])
+        self.assertIn(str(too_negative), str(cm.exception))
+
+        # Shapes that fit in int32 still go through unchanged.
+        self.assertEqual(mx.zeros(4).shape, (4,))
+        self.assertEqual(mx.zeros((2, 3)).shape, (2, 3))
+        self.assertEqual(mx.ones([2, 3]).shape, (2, 3))
+        self.assertEqual(mx.full((2, 3), 1.5).tolist(), [[1.5] * 3] * 2)
+
     def test_scalar_inputs(self):
         # Check combinations of python types
         a = mx.add(False, True)

python/tests/test_quantized.py

@@ -1046,7 +1046,7 @@ class TestQuantized(mlx_tests.MLXTestCase):
                 y3 = scatter_unsort(y3, inv_order, indices.shape)
                 y4 = scatter_unsort(y4, inv_order, indices.shape)
 
-                tol = 1.5e-5 if (dtype == mx.float32) else 2.5e-4
+                tol = 1.5e-5 if (dtype == mx.float32) else 1e-3
 
                 self.assertLess((y1 - y2).abs().max(), tol)
                 self.assertLess((y1 - y3).abs().max(), tol)

setup.py

@@ -234,7 +234,7 @@ if __name__ == "__main__":
             "ml_dtypes",
             "numpy>=2",
             "pre-commit",
-            "psutil",
+            "psutil>=7.2",
             "torch>=2.9",
             "typing_extensions",
         ],

tests/linalg_tests.cpp

@@ -637,3 +637,68 @@ TEST_CASE("test solve_triangluar") {
   expected = array({-3., 2., 3.});
   CHECK(allclose(expected, result).item<bool>());
 }
+
+TEST_CASE("test det") {
+  // 1x1 fast path
+  {
+    array a = array({5.0f}, {1, 1});
+    auto d = det(a, Device::cpu);
+    CHECK_EQ(d.item<float>(), doctest::Approx(5.0f));
+  }
+
+  // 2x2 fast path: det([[1,2],[3,4]]) = -2
+  {
+    array a = array({1.0f, 2.0f, 3.0f, 4.0f}, {2, 2});
+    auto d = det(a, Device::cpu);
+    CHECK_EQ(d.item<float>(), doctest::Approx(-2.0f));
+  }
+
+  // 3x3 fast path: det([[1,2,3],[0,1,4],[5,6,0]]) = 1
+  {
+    array a =
+        array({1.0f, 2.0f, 3.0f, 0.0f, 1.0f, 4.0f, 5.0f, 6.0f, 0.0f}, {3, 3});
+    auto d = det(a, Device::cpu);
+    CHECK_EQ(d.item<float>(), doctest::Approx(1.0f));
+  }
+
+  // 4x4 LU path: identity matrix det = 1
+  {
+    array a = eye(4);
+    auto d = det(a, Device::cpu);
+    CHECK_EQ(d.item<float>(), doctest::Approx(1.0f));
+  }
+
+  // Non-square should throw
+  CHECK_THROWS(
+      det(array({1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f}, {2, 3}), Device::cpu));
+
+  // 1D should throw
+  CHECK_THROWS(det(array({1.0f, 2.0f}), Device::cpu));
+}
+
+TEST_CASE("test slogdet") {
+  // 2x2: det = -2, so sign = -1, logabsdet = log(2)
+  {
+    array a = array({1.0f, 2.0f, 3.0f, 4.0f}, {2, 2});
+    auto [s, logabs] = slogdet(a, Device::cpu);
+    CHECK_EQ(s.item<float>(), doctest::Approx(-1.0f));
+    CHECK_EQ(logabs.item<float>(), doctest::Approx(std::log(2.0f)));
+  }
+
+  // Identity: sign = 1, logabsdet = 0
+  {
+    array a = eye(4);
+    auto [s, logabs] = slogdet(a, Device::cpu);
+    CHECK_EQ(s.item<float>(), doctest::Approx(1.0f));
+    CHECK_EQ(logabs.item<float>(), doctest::Approx(0.0f));
+  }
+
+  // Singular: sign = 0, logabsdet = -inf
+  {
+    array a = array({1.0f, 2.0f, 2.0f, 4.0f}, {2, 2});
+    auto [s, logabs] = slogdet(a, Device::cpu);
+    CHECK_EQ(s.item<float>(), 0.0f);
+    CHECK(std::isinf(logabs.item<float>()));
+    CHECK(logabs.item<float>() < 0);
+  }
+}