Bump the version (#604)

- state was being read from parameter "s"
- but being stored in parameter "u"

.circleci/config.yml

@@ -26,18 +26,23 @@ jobs:
           command: |
             pip install --upgrade cmake
             pip install --upgrade pybind11[global]
+            pip install pybind11-stubgen
             pip install numpy
             sudo apt-get update
-            sudo apt-get install libblas-dev
+            sudo apt-get install libblas-dev liblapack-dev liblapacke-dev
       - run:
-          name: Build python package
+          name: Install Python package
           command: |
             CMAKE_ARGS="-DMLX_BUILD_METAL=OFF" CMAKE_BUILD_PARALLEL_LEVEL="" python3 setup.py build_ext --inplace
             CMAKE_ARGS="-DMLX_BUILD_METAL=OFF" CMAKE_BUILD_PARALLEL_LEVEL="" python3 setup.py develop
       - run:
-          name: Run the python tests
+          name: Generate package stubs
           command: |
-            python3 -m unittest discover python/tests
+            python3 setup.py generate_stubs
+      - run:
+          name: Run Python tests
+          command: |
+            python3 -m unittest discover python/tests -v
       # TODO: Reenable when extension api becomes stable
       # - run:
       #     name: Build example extension
@@ -65,19 +70,26 @@ jobs:
             conda activate runner-env
             pip install --upgrade cmake
             pip install --upgrade pybind11[global]
+            pip install pybind11-stubgen
             pip install numpy
             pip install torch
             pip install tensorflow
             pip install unittest-xml-reporting
       - run:
-          name: Build python package
+          name: Install Python package
           command: |
             eval "$(conda shell.bash hook)"
             conda activate runner-env
             CMAKE_BUILD_PARALLEL_LEVEL="" python setup.py build_ext --inplace
             CMAKE_BUILD_PARALLEL_LEVEL="" python setup.py develop
       - run:
-          name: Run the python tests
+          name: Generate package stubs
+          command: |
+            eval "$(conda shell.bash hook)"
+            conda activate runner-env
+            python setup.py generate_stubs
+      - run:
+          name: Run Python tests
           command: |
             eval "$(conda shell.bash hook)"
             conda activate runner-env
@@ -121,10 +133,27 @@ jobs:
             conda activate runner-env
             pip install --upgrade cmake
             pip install --upgrade pybind11[global]
+            pip install pybind11-stubgen
             pip install numpy
             pip install twine
+      # TODO: Update build system to switch away from setup.py develop
       - run:
-          name: Build package
+          name: Install Python package
+          command: |
+            eval "$(conda shell.bash hook)"
+            conda activate runner-env
+            DEVELOPER_DIR=$(developer_dir_macos_<< parameters.macos_version >>) \
+              PYPI_RELEASE=1 \
+              CMAKE_BUILD_PARALLEL_LEVEL="" \
+              python setup.py develop
+      - run:
+          name: Generate package stubs
+          command: |
+            eval "$(conda shell.bash hook)"
+            conda activate runner-env
+            python setup.py generate_stubs
+      - run:
+          name: Publish Python package
           command: |
             eval "$(conda shell.bash hook)"
             conda activate runner-env
@@ -157,10 +186,26 @@ jobs:
             conda activate runner-env
             pip install --upgrade cmake
             pip install --upgrade pybind11[global]
+            pip install pybind11-stubgen
             pip install numpy
             pip install twine
       - run:
-          name: Build package
+          name: Install Python package
+          command: |
+            eval "$(conda shell.bash hook)"
+            conda activate runner-env
+            DEVELOPER_DIR=$(developer_dir_macos_<< parameters.macos_version >>) \
+              DEV_RELEASE=1 \
+              CMAKE_BUILD_PARALLEL_LEVEL="" \
+              python setup.py develop
+      - run:
+          name: Generate package stubs
+          command: |
+            eval "$(conda shell.bash hook)"
+            conda activate runner-env
+            python setup.py generate_stubs
+      - run:
+          name: Publish Python package
           command: |
             eval "$(conda shell.bash hook)"
             conda activate runner-env
@@ -193,10 +238,25 @@ jobs:
             conda activate runner-env
             pip install --upgrade cmake
             pip install --upgrade pybind11[global]
+            pip install pybind11-stubgen
             pip install numpy
             pip install twine
       - run:
-          name: Build package
+          name: Install Python package
+          command: |
+            eval "$(conda shell.bash hook)"
+            conda activate runner-env
+            DEVELOPER_DIR=$(developer_dir_macos_<< parameters.macos_version >>) \
+              CMAKE_BUILD_PARALLEL_LEVEL="" \
+              python setup.py develop
+      - run:
+          name: Generate package stubs
+          command: |
+            eval "$(conda shell.bash hook)"
+            conda activate runner-env
+            python setup.py generate_stubs
+      - run:
+          name: Build package distribution
           command: |
             eval "$(conda shell.bash hook)"
             conda activate runner-env

.pre-commit-config.yaml

@@ -9,7 +9,7 @@ repos:
     hooks:
     -   id: black
 -   repo: https://github.com/pycqa/isort
-    rev: 5.12.0
+    rev: 5.13.2
     hooks:
     -   id: isort
         args:

CMakeLists.txt

@@ -18,7 +18,7 @@ option(MLX_BUILD_METAL "Build metal backend" ON)
 option(BUILD_SHARED_LIBS "Build mlx as a shared library" OFF)
 
 if(NOT MLX_VERSION)
-  set(MLX_VERSION 0.0.10)
+  set(MLX_VERSION 0.1.0)
 endif()
 
 # --------------------- Processor tests -------------------------
@@ -31,13 +31,13 @@ if (${CMAKE_SYSTEM_NAME} MATCHES "Darwin")
 
   if (${CMAKE_HOST_SYSTEM_PROCESSOR} MATCHES "x86_64" AND ${CMAKE_HOST_APPLE})
     message(FATAL_ERROR
-      "Building for x86_64 on macOS is not supported." 
+      "Building for x86_64 on macOS is not supported."
       " If you are on an Apple silicon system, check the build"
       " documentation for possible fixes: "
       "https://ml-explore.github.io/mlx/build/html/install.html#build-from-source")
   elseif (${CMAKE_HOST_SYSTEM_PROCESSOR} MATCHES "x86_64")
-    message(WARNING 
-      "Building for x86_64 on macOS is not supported." 
+    message(WARNING
+      "Building for x86_64 on macOS is not supported."
       " If you are on an Apple silicon system, "
       " make sure you are building for arm64.")
   elseif(${CMAKE_HOST_SYSTEM_PROCESSOR} MATCHES "arm64")
@@ -75,7 +75,7 @@ elseif (MLX_BUILD_METAL)
                   COMMAND_ERROR_IS_FATAL ANY)
 
   message(STATUS "Building with SDK for macOS version ${MACOS_VERSION}")
-  
+
   if (${MACOS_VERSION} GREATER_EQUAL 14.2)
     set(METAL_CPP_URL https://developer.apple.com/metal/cpp/files/metal-cpp_macOS14.2_iOS17.2.zip)
   elseif (${MACOS_VERSION} GREATER_EQUAL 14.0)
@@ -123,16 +123,27 @@ else()
     /usr/include
     /usr/local/include
     $ENV{BLAS_HOME}/include)
-  message(STATUS ${BLAS_LIBRARIES})
-  message(STATUS ${BLAS_INCLUDE_DIRS})
+  message(STATUS "Blas lib" ${BLAS_LIBRARIES})
+  message(STATUS "Blas incclude" ${BLAS_INCLUDE_DIRS})
   target_include_directories(mlx PRIVATE ${BLAS_INCLUDE_DIRS})
   target_link_libraries(mlx ${BLAS_LIBRARIES})
+  find_package(LAPACK REQUIRED)
+  if (NOT LAPACK_FOUND)
+      message(FATAL_ERROR "Must have LAPACK installed")
+  endif()
+  find_path(LAPACK_INCLUDE_DIRS lapacke.h
+    /usr/include
+    /usr/local/include)
+  message(STATUS "Lapack lib" ${LAPACK_LIBRARIES})
+  message(STATUS "Lapack include " ${LAPACK_INCLUDE_DIRS})
+  target_include_directories(mlx PRIVATE ${LAPACK_INCLUDE_DIRS})
+  target_link_libraries(mlx ${LAPACK_LIBRARIES})
 endif()
 
 add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/mlx)
 
 target_include_directories(
-  mlx 
+  mlx
   PUBLIC
   $<BUILD_INTERFACE:${CMAKE_CURRENT_LIST_DIR}>
   $<INSTALL_INTERFACE:include>

MANIFEST.in

@@ -1,3 +1,4 @@
 include CMakeLists.txt
 recursive-include mlx/ *
 include python/src/*
+python/mlx/py.typed # support type hinting as in PEP-561

README.md

@@ -68,10 +68,18 @@ in the documentation.
 
 MLX is available on [PyPI](https://pypi.org/project/mlx/). To install the Python API, run:
 
+**With `pip`**:
+
 ```
 pip install mlx
 ```
 
+**With `conda`**:
+
+```
+conda install -c conda-forge mlx
+```
+
 Checkout the
 [documentation](https://ml-explore.github.io/mlx/build/html/install.html#)
 for more information on building the C++ and Python APIs from source.

benchmarks/python/comparative/bench_mlx.py

@@ -72,6 +72,9 @@ def _quant_matmul(x, w, s, b, transpose, group_size, bits):
 
 
 quant_matmul = {
+    "quant_matmul_32_2": partial(_quant_matmul, transpose=False, group_size=32, bits=2),
+    "quant_matmul_32_4": partial(_quant_matmul, transpose=False, group_size=32, bits=4),
+    "quant_matmul_32_8": partial(_quant_matmul, transpose=False, group_size=32, bits=8),
     "quant_matmul_64_2": partial(_quant_matmul, transpose=False, group_size=64, bits=2),
     "quant_matmul_64_4": partial(_quant_matmul, transpose=False, group_size=64, bits=4),
     "quant_matmul_64_8": partial(_quant_matmul, transpose=False, group_size=64, bits=8),
@@ -84,6 +87,15 @@ quant_matmul = {
     "quant_matmul_128_8": partial(
         _quant_matmul, transpose=False, group_size=128, bits=8
     ),
+    "quant_matmul_t_32_2": partial(
+        _quant_matmul, transpose=True, group_size=32, bits=2
+    ),
+    "quant_matmul_t_32_4": partial(
+        _quant_matmul, transpose=True, group_size=32, bits=4
+    ),
+    "quant_matmul_t_32_8": partial(
+        _quant_matmul, transpose=True, group_size=32, bits=8
+    ),
     "quant_matmul_t_64_2": partial(
         _quant_matmul, transpose=True, group_size=64, bits=2
     ),

benchmarks/python/llama_jax_bench.py

@@ -1,198 +0,0 @@
-# Copyright © 2023 Apple Inc.
-
-import math
-import time
-
-import jax
-import jax.numpy as jnp
-from flax import linen as nn
-
-
-class RoPE(nn.Module):
-    dims: int
-    traditional: bool = False
-
-    def _compute_rope(self, costheta, sintheta, x):
-        x1 = x[..., : self.dims // 2]
-        x2 = x[..., self.dims // 2 : self.dims]
-        rx1 = x1 * costheta - x2 * sintheta
-        rx2 = x1 * sintheta + x2 * costheta
-
-        if self.dims < x.shape[-1]:
-            rx = jnp.concatenate([rx1, rx2, x[..., self.dims :]], axis=-1)
-        else:
-            rx = jnp.concatenate([rx1, rx2], axis=-1)
-
-        return rx
-
-    def _compute_traditional_rope(self, costheta, sintheta, x):
-        x1 = x[..., ::2]
-        x2 = x[..., 1::2]
-        rx1 = x1 * costheta - x2 * sintheta
-        rx2 = x1 * sintheta + x2 * costheta
-
-        if self.dims < x.shape[-1]:
-            raise NotImplementedError(
-                "RoPE doesn't implement partial traditional application"
-            )
-
-        rx = jnp.concatenate([rx1[..., None], rx2[..., None]], axis=-1)
-
-        return rx
-
-    @staticmethod
-    def create_cos_sin_theta(
-        N: int,
-        D: int,
-        offset: int = 0,
-        base: float = 10000,
-        dtype=jnp.float32,
-    ):
-        D = D // 2
-        positions = jnp.arange(offset, N, dtype=dtype)
-        freqs = jnp.exp(-jnp.arange(0, D, dtype=dtype) * (math.log(base) / D))
-        theta = positions.reshape((-1, 1)) * freqs.reshape((1, -1))
-        costheta = jnp.cos(theta)
-        sintheta = jnp.sin(theta)
-
-        return costheta, sintheta
-
-    @nn.compact
-    def __call__(self, x, offset: int = 0):
-        shape = x.shape
-        x = x.reshape((-1, shape[-2], shape[-1]))
-        N = x.shape[1] + offset
-        costheta, sintheta = RoPE.create_cos_sin_theta(
-            N, self.dims, offset=offset, dtype=x.dtype
-        )
-
-        rope = (
-            self._compute_traditional_rope if self.traditional else self._compute_rope
-        )
-        rx = rope(costheta, sintheta, x)
-
-        return rx.reshape(shape)
-
-
-class LlamaAttention(nn.Module):
-    dims: int
-    num_heads: int
-    dtype: jnp.dtype
-
-    def setup(self):
-        num_heads = self.num_heads
-        dims = self.dims
-
-        self.rope = RoPE(dims // num_heads, True)
-        self.query_proj = nn.Dense(dims, use_bias=False, param_dtype=self.dtype)
-        self.key_proj = nn.Dense(dims, use_bias=False, param_dtype=self.dtype)
-        self.value_proj = nn.Dense(dims, use_bias=False, param_dtype=self.dtype)
-        self.out_proj = nn.Dense(dims, use_bias=False, param_dtype=self.dtype)
-
-    def __call__(self, queries, keys, values, mask=None, cache=None):
-        queries = self.query_proj(queries)
-        keys = self.key_proj(keys)
-        values = self.value_proj(values)
-
-        num_heads = self.num_heads
-        B, L, D = queries.shape
-        queries = queries.reshape((B, L, num_heads, -1)).transpose((0, 2, 1, 3))
-        keys = keys.reshape((B, L, num_heads, -1)).transpose((0, 2, 1, 3))
-        values = values.reshape((B, L, num_heads, -1)).transpose((0, 2, 1, 3))
-
-        if cache is not None:
-            key_cache, value_cache = cache
-            queries = self.rope(queries, offset=key_cache.shape[2])
-            keys = self.rope(keys, offset=key_cache.shape[2])
-            keys = jnp.concatenate([key_cache, keys], axis=2)
-            values = jnp.concatenate([value_cache, values], axis=2)
-        else:
-            queries = self.rope(queries)
-            keys = self.rope(keys)
-
-        # Dimensions are [batch x num heads x sequence x hidden dim]
-        scale = math.sqrt(1 / queries.shape[-1])
-        scores = (queries * scale) @ keys.transpose((0, 1, 3, 2))
-        if mask is not None:
-            scores = scores + mask
-        scores = jax.nn.softmax(scores, axis=-1)
-        values_hat = (scores @ values).transpose((0, 2, 1, 3)).reshape((B, L, -1))
-
-        return self.out_proj(values_hat), (keys, values)
-
-
-class LlamaEncoderLayer(nn.Module):
-    dims: int
-    mlp_dims: int
-    num_heads: int
-    dtype: jnp.dtype
-
-    def setup(self):
-        dims = self.dims
-        mlp_dims = self.mlp_dims
-        num_heads = self.num_heads
-
-        self.attention = LlamaAttention(dims, num_heads, dtype)
-
-        self.norm1 = nn.RMSNorm(param_dtype=self.dtype)
-        self.norm2 = nn.RMSNorm(param_dtype=self.dtype)
-
-        self.linear1 = nn.Dense(mlp_dims, use_bias=False, param_dtype=self.dtype)
-        self.linear2 = nn.Dense(mlp_dims, use_bias=False, param_dtype=self.dtype)
-        self.linear3 = nn.Dense(dims, use_bias=False, param_dtype=self.dtype)
-
-    def __call__(self, x, mask=None, cache=None):
-        y = self.norm1(x)
-        y, cache = self.attention(y, y, y, mask, cache)
-        x = x + y
-
-        y = self.norm2(x)
-        a = self.linear1(y)
-        b = self.linear2(y)
-        y = jax.nn.silu(a) * b
-        y = self.linear3(y)
-        x = x + y
-
-        return x, cache
-
-
-def measure(model, x, cache):
-    for i in range(5):
-        y, c = model(x, mask=None, cache=cache)
-        jax.block_until_ready((y, c))
-
-    start = time.time()
-    for i in range(5):
-        y, c = model(x, mask=None, cache=cache)
-        jax.block_until_ready((y, c))
-
-    end = time.time()
-    return (end - start) * 1000 / 5
-
-
-if __name__ == "__main__":
-    H = 32
-    D = 4096
-    F = 43 * 256
-    C = 1000
-    dtype = jnp.float16
-
-    k1, k2, k3, k4 = jax.random.split(jax.random.PRNGKey(0), 4)
-
-    x = jax.random.normal(k1, (1, 1, D), dtype)
-    cache = [
-        jax.random.normal(k2, [1, H, C, D // H], dtype),
-        jax.random.normal(k3, [1, H, C, D // H], dtype),
-    ]
-
-    layer = LlamaEncoderLayer(D, F, H, dtype=dtype)
-    params = layer.init(k4, x, mask=None, cache=cache)["params"]
-
-    @jax.jit
-    def model_fn(x, mask, cache):
-        return layer.apply({"params": params}, x, mask=mask, cache=cache)
-
-    T = measure(model_fn, x, cache)
-
-    print("Time per layer per token:", T, "ms")
-    print("Lower bound total time per token:", T * 32, "ms")

benchmarks/python/llama_mlx_bench.py

@@ -1,118 +0,0 @@
-# Copyright © 2023 Apple Inc.
-
-import math
-import time
-
-import mlx.core as mx
-import mlx.nn as nn
-import mlx.utils
-
-
-class LlamaAttention(nn.Module):
-    def __init__(self, dims: int, num_heads: int):
-        super().__init__()
-        self.num_heads = num_heads
-        self.rope = nn.RoPE(dims // num_heads, True)
-        self.query_proj = nn.Linear(dims, dims, False)
-        self.key_proj = nn.Linear(dims, dims, False)
-        self.value_proj = nn.Linear(dims, dims, False)
-        self.out_proj = nn.Linear(dims, dims, False)
-
-    def __call__(self, queries, keys, values, mask=None, cache=None):
-        queries = self.query_proj(queries)
-        keys = self.key_proj(keys)
-        values = self.value_proj(values)
-
-        num_heads = self.num_heads
-        B, L, D = queries.shape
-        queries = mx.transpose(mx.reshape(queries, (B, L, num_heads, -1)), (0, 2, 1, 3))
-        keys = mx.transpose(mx.reshape(keys, (B, L, num_heads, -1)), (0, 2, 1, 3))
-        values = mx.transpose(mx.reshape(values, (B, L, num_heads, -1)), (0, 2, 1, 3))
-
-        if cache is not None:
-            key_cache, value_cache = cache
-            queries = self.rope(queries, offset=key_cache.shape[2])
-            keys = self.rope(keys, offset=key_cache.shape[2])
-            keys = mx.concatenate([key_cache, keys], axis=2)
-            values = mx.concatenate([value_cache, values], axis=2)
-        else:
-            queries = self.rope(queries)
-            keys = self.rope(keys)
-
-        # Dimensions are [batch x num heads x sequence x hidden dim]
-        scale = mx.array(math.sqrt(1 / queries.shape[-1]), dtype=queries.dtype)
-        scores = (queries * scale) @ mx.transpose(keys, (0, 1, 3, 2))
-        if mask is not None:
-            scores = scores + mask
-        scores = mx.softmax(scores, axis=-1)
-        values_hat = mx.reshape(mx.transpose(scores @ values, (0, 2, 1, 3)), (B, L, -1))
-
-        return self.out_proj(values_hat), (keys, values)
-
-
-class LlamaEncoderLayer(nn.Module):
-    def __init__(self, dims: int, mlp_dims: int, num_heads: int):
-        super().__init__()
-
-        self.attention = LlamaAttention(dims, num_heads)
-
-        self.norm1 = nn.RMSNorm(dims)
-        self.norm2 = nn.RMSNorm(dims)
-
-        self.linear1 = nn.Linear(dims, mlp_dims, False)
-        self.linear2 = nn.Linear(dims, mlp_dims, False)
-        self.linear3 = nn.Linear(mlp_dims, dims, False)
-
-    def __call__(self, x, mask=None, cache=None):
-        y = self.norm1(x)
-        y, cache = self.attention(y, y, y, mask, cache)
-        x = x + y
-
-        y = self.norm2(x)
-        a = self.linear1(y)
-        b = self.linear2(y)
-        y = a * mx.sigmoid(a) * b
-        y = self.linear3(y)
-        x = x + y
-
-        return x, cache
-
-
-def measure(model, x, cache):
-    for i in range(5):
-        y, c = model(x, mask=None, cache=cache)
-        mx.eval(y, c)
-
-    start = time.time()
-    rs = []
-    for i in range(5):
-        y, c = model(x, mask=None, cache=cache)
-        rs.append((y, c))
-    mx.eval(rs)
-    end = time.time()
-
-    return (end - start) * 1000 / 5
-
-
-if __name__ == "__main__":
-    H = 32
-    D = 4096
-    F = 43 * 256
-    C = 1000
-    mx.set_default_device(mx.gpu)
-    dtype = mx.float16
-
-    layer = LlamaEncoderLayer(D, F, H)
-    layer.update(mlx.utils.tree_map(lambda x: x.astype(dtype), layer.parameters()))
-    k1, k2, k3 = mx.random.split(mx.random.key(0), 3)
-    x = mx.random.normal([1, 1, D], dtype=dtype)
-    cache = [
-        mx.random.normal([1, H, C, D // H], dtype=dtype),
-        mx.random.normal([1, H, C, D // H], dtype=dtype),
-    ]
-    mx.eval(x, cache)
-
-    T = measure(layer, x, cache)
-
-    print("Time per layer per token:", T, "ms")
-    print("Lower bound total time per token:", T * 32, "ms")

benchmarks/python/llama_torch_bench.py

@@ -1,199 +0,0 @@
-# Copyright © 2023 Apple Inc.
-
-import math
-import time
-
-import torch
-import torch.mps
-import torch.nn as nn
-
-
-def sync_if_needed(x):
-    if x.device != torch.device("cpu"):
-        torch.mps.synchronize()
-
-
-class RoPE(nn.Module):
-    def __init__(self, dims: int, traditional: bool = False):
-        super().__init__()
-        self.dims = dims
-        self.traditional = traditional
-
-    def _compute_rope(self, costheta, sintheta, x):
-        x1 = x[..., : self.dims // 2]
-        x2 = x[..., self.dims // 2 : self.dims]
-        rx1 = x1 * costheta - x2 * sintheta
-        rx2 = x1 * sintheta + x2 * costheta
-
-        if self.dims < x.shape[-1]:
-            rx = torch.cat([rx1, rx2, x[..., self.dims :]], dim=-1)
-        else:
-            rx = torch.cat([rx1, rx2], dim=-1)
-
-        return rx
-
-    def _compute_traditional_rope(self, costheta, sintheta, x):
-        x1 = x[..., ::2]
-        x2 = x[..., 1::2]
-        rx1 = x1 * costheta - x2 * sintheta
-        rx2 = x1 * sintheta + x2 * costheta
-
-        if self.dims < x.shape[-1]:
-            raise NotImplementedError(
-                "RoPE doesn't implement partial traditional application"
-            )
-
-        rx = torch.cat([rx1[..., None], rx2[..., None]], dim=-1)
-
-        return rx
-
-    def forward(self, x, offset: int = 0):
-        shape = x.shape
-        x = x.view(-1, shape[-2], shape[-1])
-        N = x.shape[1] + offset
-        costheta, sintheta = RoPE.create_cos_sin_theta(
-            N, self.dims, offset=offset, device=x.device, dtype=x.dtype
-        )
-
-        rope = (
-            self._compute_traditional_rope if self.traditional else self._compute_rope
-        )
-        rx = rope(costheta, sintheta, x)
-
-        return rx.view(*shape)
-
-    @staticmethod
-    def create_cos_sin_theta(
-        N: int,
-        D: int,
-        offset: int = 0,
-        base: float = 10000,
-        device="cpu",
-        dtype=torch.float32,
-    ):
-        D = D // 2
-        positions = torch.arange(offset, N, dtype=dtype, device=device)
-        freqs = torch.exp(
-            -torch.arange(0, D, dtype=dtype, device=device) * (math.log(base) / D)
-        )
-        theta = positions.view(-1, 1) * freqs.view(1, -1)
-        costheta = torch.cos(theta)
-        sintheta = torch.sin(theta)
-
-        return costheta, sintheta
-
-
-class RMSNorm(nn.Module):
-    def __init__(self, dims: int, epsilon: float = 1e-6):
-        super().__init__()
-        self.gamma = nn.Parameter(torch.ones((dims,)))
-        self.epsilon = epsilon
-
-    def forward(self, x):
-        n = torch.rsqrt(x.square().mean(dim=-1, keepdims=True) + self.epsilon)
-        return self.gamma * x * n
-
-
-class LlamaAttention(nn.Module):
-    def __init__(self, dims: int, num_heads: int):
-        super().__init__()
-        self.num_heads = num_heads
-        self.rope = RoPE(dims // num_heads, True)
-        self.query_proj = nn.Linear(dims, dims, bias=False)
-        self.key_proj = nn.Linear(dims, dims, bias=False)
-        self.value_proj = nn.Linear(dims, dims, bias=False)
-        self.out_proj = nn.Linear(dims, dims, bias=False)
-
-    def forward(self, queries, keys, values, mask=None, cache=None):
-        queries = self.query_proj(queries)
-        keys = self.key_proj(keys)
-        values = self.value_proj(values)
-
-        num_heads = self.num_heads
-        B, L, D = queries.shape
-        queries = queries.view(B, L, num_heads, -1).permute(0, 2, 1, 3)
-        keys = keys.view(B, L, num_heads, -1).permute(0, 2, 1, 3)
-        values = values.view(B, L, num_heads, -1).permute(0, 2, 1, 3)
-
-        if cache is not None:
-            key_cache, value_cache = cache
-            queries = self.rope(queries, offset=key_cache.shape[2])
-            keys = self.rope(keys, offset=key_cache.shape[2])
-            keys = torch.cat([key_cache, keys], dim=2)
-            values = torch.cat([value_cache, values], dim=2)
-        else:
-            queries = self.rope(queries)
-            keys = self.rope(keys)
-
-        # Dimensions are [batch x num heads x sequence x hidden dim]
-        scale = math.sqrt(1 / queries.shape[-1])
-        scores = (queries * scale) @ keys.permute(0, 1, 3, 2)
-        if mask is not None:
-            scores = scores + mask
-        scores = torch.softmax(scores, dim=-1)
-        values_hat = (scores @ values).permute(0, 2, 1, 3).reshape(B, L, -1)
-
-        return self.out_proj(values_hat), (keys, values)
-
-
-class LlamaEncoderLayer(nn.Module):
-    def __init__(self, dims: int, mlp_dims: int, num_heads: int):
-        super().__init__()
-
-        self.attention = LlamaAttention(dims, num_heads)
-
-        self.norm1 = RMSNorm(dims)
-        self.norm2 = RMSNorm(dims)
-
-        self.linear1 = nn.Linear(dims, mlp_dims, bias=False)
-        self.linear2 = nn.Linear(dims, mlp_dims, bias=False)
-        self.linear3 = nn.Linear(mlp_dims, dims, bias=False)
-
-    def forward(self, x, mask=None, cache=None):
-        y = self.norm1(x)
-        y, cache = self.attention(y, y, y, mask, cache)
-        x = x + y
-
-        y = self.norm2(x)
-        a = self.linear1(y)
-        b = self.linear2(y)
-        y = torch.nn.functional.silu(a) * b
-        y = self.linear3(y)
-        x = x + y
-
-        return x, cache
-
-
-@torch.no_grad()
-def measure(model, x, cache):
-    for i in range(5):
-        y, c = model(x, mask=None, cache=cache)
-    sync_if_needed(x)
-
-    start = time.time()
-    for i in range(5):
-        y, c = model(x, mask=None, cache=cache)
-    sync_if_needed(x)
-    end = time.time()
-    return (end - start) * 1000 / 5
-
-
-if __name__ == "__main__":
-    H = 32
-    D = 4096
-    F = 43 * 256
-    C = 1000
-    device = torch.device("mps")
-    dtype = torch.float16
-
-    layer = LlamaEncoderLayer(D, F, H).to(device).to(dtype)
-    x = torch.randn(1, 1, D).to(device).to(dtype)
-    cache = [
-        torch.randn(1, H, C, D // H).to(device).to(dtype),
-        torch.randn(1, H, C, D // H).to(device).to(dtype),
-    ]
-
-    T = measure(layer, x, cache)
-
-    print("Time per layer per token:", T, "ms")
-    print("Lower bound total time per token:", T * 32, "ms")

benchmarks/python/single_ops.py

@@ -44,6 +44,13 @@ def time_matmul():
     time_fn(mx.matmul, a, b)
 
 
+def time_maximum():
+    a = mx.random.uniform(shape=(32, 1024, 1024))
+    b = mx.random.uniform(shape=(32, 1024, 1024))
+    mx.eval(a, b)
+    time_fn(mx.maximum, a, b)
+
+
 def time_negative():
     a = mx.random.uniform(shape=(10000, 1000))
     mx.eval(a)
@@ -101,6 +108,7 @@ if __name__ == "__main__":
 
     time_add()
     time_matmul()
+    time_maximum()
     time_exp()
     time_negative()
     time_logsumexp()

docs/src/conf.py

@@ -12,7 +12,7 @@ import mlx.core as mx
 project = "MLX"
 copyright = "2023, MLX Contributors"
 author = "MLX Contributors"
-version = ".".join(mx.__version__.split()[:-1])
+version = ".".join(mx.__version__.split(".")[:3])
 release = version
 
 # -- General configuration ---------------------------------------------------

docs/src/dev/extensions.rst

@@ -929,7 +929,7 @@ We see some modest improvements right away!
 
 This operation is now good to be used to build other operations, 
 in :class:`mlx.nn.Module` calls, and also as a part of graph 
-transformations such as :meth:`grad` and :meth:`simplify`!
+transformations like :meth:`grad`!
 
 Scripts
 -------

docs/src/python/linalg.rst

@@ -9,3 +9,4 @@ Linear Algebra
    :toctree: _autosummary 
 
     norm
+    qr

docs/src/python/nn.rst

@@ -180,3 +180,4 @@ In detail:
    nn/layers
    nn/functions
    nn/losses
+   nn/init

docs/src/python/nn/functions.rst

@@ -19,5 +19,6 @@ simple functions.
    prelu
    relu
    selu
+   softshrink
    silu
    step

docs/src/python/nn/init.rst

@@ -0,0 +1,45 @@
+.. _init:
+
+.. currentmodule:: mlx.nn.init
+
+Initializers
+------------
+
+The ``mlx.nn.init`` package contains commonly used initializers for neural
+network parameters. Initializers return a function which can be applied to any
+input :obj:`mlx.core.array` to produce an initialized output.
+
+For example:
+
+.. code:: python
+
+   import mlx.core as mx
+   import mlx.nn as nn
+
+   init_fn = nn.init.uniform()
+
+   # Produces a [2, 2] uniform matrix
+   param = init_fn(mx.zeros((2, 2)))
+
+To re-initialize all the parameter in an :obj:`mlx.nn.Module` from say a uniform 
+distribution, you can do:
+
+.. code:: python
+  
+   import mlx.nn as nn
+   model = nn.Sequential(nn.Linear(5, 10), nn.ReLU(), nn.Linear(10, 5))
+   init_fn = nn.init.uniform(low=-0.1, high=0.1)
+   model.apply(init_fn)
+   
+
+.. autosummary::
+   :toctree: _autosummary
+
+   constant
+   normal
+   uniform
+   identity
+   glorot_normal
+   glorot_uniform
+   he_normal
+   he_uniform

docs/src/python/nn/layers.rst

@@ -33,5 +33,6 @@ Layers
    Sequential
    SiLU
    SinusoidalPositionalEncoding
+   Softshrink
    Step
    Transformer

docs/src/python/ops.rst

@@ -35,6 +35,8 @@ Operations
    cos
    cosh
    dequantize
+   diag
+   diagonal
    divide
    divmod
    equal

docs/src/python/optimizers.rst

@@ -40,6 +40,7 @@ model's parameters and the **optimizer state**.
    SGD
    RMSprop
    Adagrad
+   Adafactor
    AdaDelta
    Adam
    AdamW

docs/src/python/transforms.rst

@@ -14,4 +14,3 @@ Transforms
    jvp
    vjp
    vmap
-   simplify

docs/src/usage/lazy_evaluation.rst

@@ -20,7 +20,7 @@ Transforming Compute Graphs
 
 Lazy evaluation let's us record a compute graph without actually doing any
 computations. This is useful for function transformations like :func:`grad` and
-:func:`vmap` and graph optimizations like :func:`simplify`.
+:func:`vmap` and graph optimizations.
 
 Currently, MLX does not compile and rerun compute graphs. They are all
 generated dynamically. However, lazy evaluation makes it much easier to

examples/python/linear_regression.py

@@ -41,6 +41,6 @@ error_norm = mx.sum(mx.square(w - w_star)).item() ** 0.5
 throughput = num_iters / (toc - tic)
 
 print(
-    f"Loss {loss.item():.5f}, |w-w*| = {error_norm:.5f}, "
+    f"Loss {loss.item():.5f}, L2 distance: |w-w*| = {error_norm:.5f}, "
     f"Throughput {throughput:.5f} (it/s)"
 )

mlx/CMakeLists.txt

@@ -5,6 +5,7 @@ target_sources(
   ${CMAKE_CURRENT_SOURCE_DIR}/array.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/device.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/dtype.cpp
+  ${CMAKE_CURRENT_SOURCE_DIR}/compile.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/fft.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/ops.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/graph_utils.cpp
@@ -19,7 +20,7 @@ target_sources(
 
 add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backend/common)
 add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/io)
-if (MLX_BUILD_ACCELERATE) 
+if (MLX_BUILD_ACCELERATE)
   add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backend/accelerate)
 else()
   target_sources(

mlx/array.cpp

@@ -1,4 +1,4 @@
-// Copyright © 2023 Apple Inc.
+// Copyright © 2023-2024 Apple Inc.
 
 #include <functional>
 
@@ -47,6 +47,17 @@ array::array(
           std::move(primitive),
           inputs)) {}
 
+array::array(
+    std::vector<int> shape,
+    Dtype dtype,
+    std::shared_ptr<Primitive> primitive,
+    std::vector<array>&& inputs)
+    : array_desc_(std::make_shared<ArrayDesc>(
+          std::move(shape),
+          dtype,
+          std::move(primitive),
+          std::move(inputs))) {}
+
 std::vector<array> array::make_arrays(
     const std::vector<std::vector<int>>& shapes,
     const std::vector<Dtype>& dtypes,
@@ -86,11 +97,13 @@ void array::detach() {
     s.array_desc_->inputs.clear();
     s.array_desc_->siblings.clear();
     s.array_desc_->position = 0;
+    s.array_desc_->depth = 0;
     s.array_desc_->primitive = nullptr;
   }
   array_desc_->inputs.clear();
   array_desc_->siblings.clear();
   array_desc_->position = 0;
+  array_desc_->depth = 0;
   array_desc_->primitive = nullptr;
 }
 
@@ -144,6 +157,14 @@ void array::copy_shared_buffer(const array& other) {
   copy_shared_buffer(other, other.strides(), other.flags(), other.data_size());
 }
 
+void array::move_shared_buffer(array other) {
+  array_desc_->data = std::move(other.array_desc_->data);
+  array_desc_->strides = other.strides();
+  array_desc_->flags = other.flags();
+  array_desc_->data_size = other.data_size();
+  array_desc_->data_ptr = other.array_desc_->data_ptr;
+}
+
 array::ArrayDesc::ArrayDesc(const std::vector<int>& shape, Dtype dtype)
     : shape(shape), dtype(dtype) {
   std::tie(size, strides) = cum_prod(shape);
@@ -158,10 +179,29 @@ array::ArrayDesc::ArrayDesc(
       dtype(dtype),
       primitive(std::move(primitive)),
       inputs(inputs) {
-  std::tie(size, strides) = cum_prod(shape);
+  std::tie(size, strides) = cum_prod(this->shape);
   for (auto& in : inputs) {
     is_tracer |= in.is_tracer();
+    depth = std::max(in.graph_depth(), depth);
   }
+  depth++;
+}
+
+array::ArrayDesc::ArrayDesc(
+    std::vector<int>&& shape,
+    Dtype dtype,
+    std::shared_ptr<Primitive> primitive,
+    std::vector<array>&& inputs)
+    : shape(std::move(shape)),
+      dtype(dtype),
+      primitive(std::move(primitive)),
+      inputs(std::move(inputs)) {
+  std::tie(size, strides) = cum_prod(this->shape);
+  for (auto& in : inputs) {
+    is_tracer |= in.is_tracer();
+    depth = std::max(in.graph_depth(), depth);
+  }
+  depth++;
 }
 
 array::ArrayIterator::ArrayIterator(const array& arr, int idx)

mlx/array.h

@@ -172,6 +172,12 @@ class array {
       std::shared_ptr<Primitive> primitive,
       const std::vector<array>& inputs);
 
+  array(
+      std::vector<int> shape,
+      Dtype dtype,
+      std::shared_ptr<Primitive> primitive,
+      std::vector<array>&& inputs);
+
   static std::vector<array> make_arrays(
       const std::vector<std::vector<int>>& shapes,
       const std::vector<Dtype>& dtypes,
@@ -215,6 +221,11 @@ class array {
     return *(array_desc_->primitive);
   };
 
+  /** A shared pointer to the array's primitive. */
+  std::shared_ptr<Primitive>& primitive_ptr() const {
+    return array_desc_->primitive;
+  };
+
   /** Check if the array has an attached primitive or is a leaf node. */
   bool has_primitive() const {
     return array_desc_->primitive != nullptr;
@@ -229,6 +240,11 @@ class array {
     return array_desc_->inputs;
   }
 
+  /** True indicates the arrays buffer is safe to reuse */
+  bool is_donatable() const {
+    return array_desc_.use_count() == 1 && (array_desc_->data.use_count() == 1);
+  }
+
   /** The array's siblings. */
   const std::vector<array>& siblings() const {
     return array_desc_->siblings;
@@ -251,6 +267,11 @@ class array {
     return outputs;
   };
 
+  /** The depth of the array in the graph. Evaluated arrays have depth 0. */
+  uint16_t graph_depth() const {
+    return array_desc_->depth;
+  }
+
   /** Detach the array from the graph. */
   void detach();
 
@@ -271,6 +292,12 @@ class array {
     return array_desc_->data->buffer;
   };
 
+  // Return a copy of the shared pointer
+  // to the array::Data struct
+  std::shared_ptr<Data> data_shared_ptr() const {
+    return array_desc_->data;
+  }
+  // Return a raw pointer to the arrays data
   template <typename T>
   T* data() {
     return static_cast<T*>(array_desc_->data_ptr);
@@ -311,6 +338,8 @@ class array {
 
   void copy_shared_buffer(const array& other);
 
+  void move_shared_buffer(array other);
+
   void overwrite_descriptor(const array& other) {
     array_desc_ = other.array_desc_;
   }
@@ -353,6 +382,9 @@ class array {
     // The arrays position in the output list
     uint32_t position{0};
 
+    // The depth of the array in the graph.
+    uint16_t depth{0};
+
     explicit ArrayDesc(const std::vector<int>& shape, Dtype dtype);
 
     explicit ArrayDesc(
@@ -360,12 +392,18 @@ class array {
         Dtype dtype,
         std::shared_ptr<Primitive> primitive,
         const std::vector<array>& inputs);
+
+    explicit ArrayDesc(
+        std::vector<int>&& shape,
+        Dtype dtype,
+        std::shared_ptr<Primitive> primitive,
+        std::vector<array>&& inputs);
   };
 
   // The ArrayDesc contains the details of the materialized array including the
   // shape, strides, the data type. It also includes
   // the primitive which knows how to compute the array's data from its inputs
-  // and a the list of array's inputs for the primitive.
+  // and the list of array's inputs for the primitive.
   std::shared_ptr<ArrayDesc> array_desc_{nullptr};
 };
 

mlx/backend/accelerate/matmul.cpp

@@ -46,6 +46,11 @@ inline void matmul_cblas_general(
   size_t N = b.shape(-1);
   size_t K = a.shape(-1);
 
+  if (K == 0) {
+    std::memset(static_cast<void*>(out.data<float>()), 0, out.nbytes());
+    return;
+  }
+
   for (int i = 0; i < (a.size() / (M * K)); ++i) {
     cblas_sgemm(
         CblasRowMajor,
@@ -89,6 +94,11 @@ inline void matmul_bnns_general(
   size_t N = b.shape(-1);
   size_t K = a.shape(-1);
 
+  if (K == 0) {
+    std::memset(static_cast<void*>(out.data<float>()), 0, out.nbytes());
+    return;
+  }
+
   BNNSDataType bnns_dtype = to_bnns_dtype(out.dtype());
 
   const BNNSLayerParametersBroadcastMatMul gemm_params{
@@ -201,4 +211,4 @@ void AddMM::eval_cpu(const std::vector<array>& inputs, array& out) {
   return matmul_bnns_general(inputs[0], inputs[1], out, alpha_, beta_);
 }
 
-} // namespace mlx::core
+} // namespace mlx::core

mlx/backend/accelerate/primitives.cpp

@@ -1,4 +1,4 @@
-// Copyright © 2023 Apple Inc.
+// Copyright © 2023-2024 Apple Inc.
 
 #include <cassert>
 #include <cmath>
@@ -35,6 +35,8 @@ DEFAULT(Broadcast)
 DEFAULT(Ceil)
 DEFAULT(Concatenate)
 DEFAULT(Copy)
+DEFAULT_MULTI(CustomVJP)
+DEFAULT_MULTI(Depends)
 DEFAULT(Equal)
 DEFAULT(Erf)
 DEFAULT(ErfInv)
@@ -50,6 +52,8 @@ DEFAULT(LogicalNot)
 DEFAULT(LogicalAnd)
 DEFAULT(LogicalOr)
 DEFAULT(LogAddExp)
+DEFAULT(Maximum)
+DEFAULT(Minimum)
 DEFAULT(NotEqual)
 DEFAULT(Pad)
 DEFAULT(Partition)
@@ -65,26 +69,17 @@ DEFAULT(Sort)
 DEFAULT(StopGradient)
 DEFAULT(Transpose)
 DEFAULT_MULTI(DivMod)
+DEFAULT_MULTI(QRF)
 
 void Abs::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   auto& in = inputs[0];
   if (in.dtype() == float32 && in.flags().contiguous) {
-    auto size = in.data_size();
-    out.set_data(
-        allocator::malloc_or_wait(size * out.itemsize()),
-        size,
-        in.strides(),
-        in.flags());
-    vDSP_vabs(in.data<float>(), 1, out.data<float>(), 1, size);
+    set_unary_output_data(in, out);
+    vDSP_vabs(in.data<float>(), 1, out.data<float>(), 1, in.data_size());
   } else if (in.dtype() == int32 && in.flags().contiguous) {
-    auto size = in.data_size();
-    out.set_data(
-        allocator::malloc_or_wait(size * out.itemsize()),
-        size,
-        in.strides(),
-        in.flags());
-    vDSP_vabsi(in.data<int>(), 1, out.data<int>(), 1, size);
+    set_unary_output_data(in, out);
+    vDSP_vabsi(in.data<int>(), 1, out.data<int>(), 1, in.data_size());
   } else if (is_unsigned(in.dtype())) {
     // No-op for unsigned types
     out.copy_shared_buffer(in);
@@ -137,12 +132,8 @@ void ArcCos::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   const auto& in = inputs[0];
   if (out.dtype() == float32 && in.flags().contiguous) {
+    set_unary_output_data(in, out);
     int size = in.data_size();
-    out.set_data(
-        allocator::malloc_or_wait(size * out.itemsize()),
-        size,
-        in.strides(),
-        in.flags());
     vvacosf(out.data<float>(), in.data<float>(), &size);
   } else {
     eval(inputs, out);
@@ -153,12 +144,8 @@ void ArcCosh::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   const auto& in = inputs[0];
   if (out.dtype() == float32 && in.flags().contiguous) {
+    set_unary_output_data(in, out);
     int size = in.data_size();
-    out.set_data(
-        allocator::malloc_or_wait(size * out.itemsize()),
-        size,
-        in.strides(),
-        in.flags());
     vvacoshf(out.data<float>(), in.data<float>(), &size);
   } else {
     eval(inputs, out);
@@ -169,12 +156,8 @@ void ArcSin::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   const auto& in = inputs[0];
   if (out.dtype() == float32 && in.flags().contiguous) {
+    set_unary_output_data(in, out);
     int size = in.data_size();
-    out.set_data(
-        allocator::malloc_or_wait(size * out.itemsize()),
-        size,
-        in.strides(),
-        in.flags());
     vvasinf(out.data<float>(), in.data<float>(), &size);
   } else {
     eval(inputs, out);
@@ -185,12 +168,8 @@ void ArcSinh::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   const auto& in = inputs[0];
   if (out.dtype() == float32 && in.flags().contiguous) {
+    set_unary_output_data(in, out);
     int size = in.data_size();
-    out.set_data(
-        allocator::malloc_or_wait(size * out.itemsize()),
-        size,
-        in.strides(),
-        in.flags());
     vvasinhf(out.data<float>(), in.data<float>(), &size);
   } else {
     eval(inputs, out);
@@ -201,12 +180,8 @@ void ArcTan::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   const auto& in = inputs[0];
   if (out.dtype() == float32 && in.flags().contiguous) {
+    set_unary_output_data(in, out);
     int size = in.data_size();
-    out.set_data(
-        allocator::malloc_or_wait(size * out.itemsize()),
-        size,
-        in.strides(),
-        in.flags());
     vvatanf(out.data<float>(), in.data<float>(), &size);
   } else {
     eval(inputs, out);
@@ -217,12 +192,8 @@ void ArcTanh::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   const auto& in = inputs[0];
   if (out.dtype() == float32 && in.flags().contiguous) {
+    set_unary_output_data(in, out);
     int size = in.data_size();
-    out.set_data(
-        allocator::malloc_or_wait(size * out.itemsize()),
-        size,
-        in.strides(),
-        in.flags());
     vvatanhf(out.data<float>(), in.data<float>(), &size);
   } else {
     eval(inputs, out);
@@ -234,30 +205,23 @@ void AsType::eval_cpu(const std::vector<array>& inputs, array& out) {
   auto& in = inputs[0];
 
   if (in.flags().contiguous) {
-    auto allocfn = [&in, &out]() {
-      out.set_data(
-          allocator::malloc_or_wait(in.data_size() * out.itemsize()),
-          in.data_size(),
-          in.strides(),
-          in.flags());
-    };
     // Use accelerate functions if possible
     if (in.dtype() == float32 && out.dtype() == uint32) {
-      allocfn();
+      set_unary_output_data(in, out);
       vDSP_vfixu32(
           in.data<float>(), 1, out.data<uint32_t>(), 1, in.data_size());
       return;
     } else if (in.dtype() == float32 && out.dtype() == int32) {
-      allocfn();
+      set_unary_output_data(in, out);
       vDSP_vfix32(in.data<float>(), 1, out.data<int32_t>(), 1, in.data_size());
       return;
     } else if (in.dtype() == uint32 && out.dtype() == float32) {
-      allocfn();
+      set_unary_output_data(in, out);
       vDSP_vfltu32(
           in.data<uint32_t>(), 1, out.data<float>(), 1, in.data_size());
       return;
     } else if (in.dtype() == int32 && out.dtype() == float32) {
-      allocfn();
+      set_unary_output_data(in, out);
       vDSP_vflt32(in.data<int32_t>(), 1, out.data<float>(), 1, in.data_size());
       return;
     }
@@ -269,12 +233,8 @@ void Cos::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   const auto& in = inputs[0];
   if (out.dtype() == float32 && in.flags().contiguous) {
+    set_unary_output_data(in, out);
     int size = in.data_size();
-    out.set_data(
-        allocator::malloc_or_wait(size * out.itemsize()),
-        size,
-        in.strides(),
-        in.flags());
     vvcosf(out.data<float>(), in.data<float>(), &size);
   } else {
     eval(inputs, out);
@@ -285,12 +245,8 @@ void Cosh::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   const auto& in = inputs[0];
   if (out.dtype() == float32 && in.flags().contiguous) {
+    set_unary_output_data(in, out);
     int size = in.data_size();
-    out.set_data(
-        allocator::malloc_or_wait(size * out.itemsize()),
-        size,
-        in.strides(),
-        in.flags());
     vvcoshf(out.data<float>(), in.data<float>(), &size);
   } else {
     eval(inputs, out);
@@ -378,12 +334,8 @@ void Exp::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   const auto& in = inputs[0];
   if (out.dtype() == float32 && in.flags().contiguous) {
+    set_unary_output_data(in, out);
     auto size = in.data_size();
-    out.set_data(
-        allocator::malloc_or_wait(size * out.itemsize()),
-        size,
-        in.strides(),
-        in.flags());
     vvexpf(out.data<float>(), in.data<float>(), reinterpret_cast<int*>(&size));
   } else if (is_floating_point(out.dtype())) {
     unary_fp(in, out, [](auto x) { return std::exp(x); });
@@ -410,12 +362,8 @@ void Log::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   const auto& in = inputs[0];
   if (out.dtype() == float32 && in.flags().contiguous) {
+    set_unary_output_data(in, out);
     auto size = in.data_size();
-    out.set_data(
-        allocator::malloc_or_wait(size * out.itemsize()),
-        size,
-        in.strides(),
-        in.flags());
     switch (base_) {
       case Base::e:
         vvlogf(
@@ -439,12 +387,8 @@ void Log1p::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   const auto& in = inputs[0];
   if (out.dtype() == float32 && in.flags().contiguous) {
+    set_unary_output_data(in, out);
     auto size = in.data_size();
-    out.set_data(
-        allocator::malloc_or_wait(size * out.itemsize()),
-        size,
-        in.strides(),
-        in.flags());
     vvlog1pf(
         out.data<float>(), in.data<float>(), reinterpret_cast<int*>(&size));
   } else if (is_floating_point(out.dtype())) {
@@ -456,47 +400,6 @@ void Log1p::eval_cpu(const std::vector<array>& inputs, array& out) {
   }
 }
 
-void Maximum::eval_cpu(const std::vector<array>& inputs, array& out) {
-  assert(inputs.size() == 2);
-  auto& a = inputs[0];
-  auto& b = inputs[1];
-  if (out.dtype() == float32) {
-    binary(
-        a,
-        b,
-        out,
-        [](auto x, auto y) { return (x > y) ? x : y; },
-        UseDefaultBinaryOp(),
-        UseDefaultBinaryOp(),
-        [](const auto* a, const auto* b, auto* out, int n) {
-          vDSP_vmax((const float*)a, 1, (const float*)b, 1, (float*)out, 1, n);
-        });
-  } else {
-    binary(a, b, out, [](auto x, auto y) { return (x > y) ? x : y; });
-  }
-}
-
-void Minimum::eval_cpu(const std::vector<array>& inputs, array& out) {
-  assert(inputs.size() == 2);
-  auto& a = inputs[0];
-  auto& b = inputs[1];
-
-  if (out.dtype() == float32) {
-    binary(
-        a,
-        b,
-        out,
-        [](auto x, auto y) { return (x < y) ? x : y; },
-        UseDefaultBinaryOp(),
-        UseDefaultBinaryOp(),
-        [](const auto* a, const auto* b, auto* out, int n) {
-          vDSP_vmin((const float*)a, 1, (const float*)b, 1, (float*)out, 1, n);
-        });
-  } else {
-    binary(a, b, out, [](auto x, auto y) { return (x < y) ? x : y; });
-  }
-}
-
 void Multiply::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 2);
   auto& a = inputs[0];
@@ -526,13 +429,8 @@ void Negative::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   auto& in = inputs[0];
   if (in.dtype() == float32 && in.flags().contiguous) {
-    auto size = in.data_size();
-    out.set_data(
-        allocator::malloc_or_wait(size * out.itemsize()),
-        size,
-        in.strides(),
-        in.flags());
-    vDSP_vneg(in.data<float>(), 1, out.data<float>(), 1, size);
+    set_unary_output_data(in, out);
+    vDSP_vneg(in.data<float>(), 1, out.data<float>(), 1, in.data_size());
   } else {
     unary(in, out, [](auto x) { return -x; });
   }
@@ -545,7 +443,13 @@ void Power::eval_cpu(const std::vector<array>& inputs, array& out) {
   if (out.dtype() == float32 && a.flags().row_contiguous &&
       b.flags().row_contiguous) {
     int size = a.size();
-    out.set_data(allocator::malloc_or_wait(out.nbytes()));
+    if (a.is_donatable() && a.itemsize() == out.itemsize()) {
+      out.copy_shared_buffer(a);
+    } else if (b.is_donatable() && b.itemsize() == out.itemsize()) {
+      out.copy_shared_buffer(b);
+    } else {
+      out.set_data(allocator::malloc_or_wait(out.nbytes()));
+    }
     vvpowf(out.data<float>(), b.data<float>(), a.data<float>(), &size);
   } else {
     eval(inputs, out);
@@ -587,12 +491,8 @@ void Sin::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   const auto& in = inputs[0];
   if (out.dtype() == float32 && in.flags().contiguous) {
+    set_unary_output_data(in, out);
     int size = in.data_size();
-    out.set_data(
-        allocator::malloc_or_wait(size * out.itemsize()),
-        size,
-        in.strides(),
-        in.flags());
     vvsinf(out.data<float>(), in.data<float>(), &size);
   } else {
     eval(inputs, out);
@@ -603,12 +503,8 @@ void Sinh::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   const auto& in = inputs[0];
   if (out.dtype() == float32 && in.flags().contiguous) {
+    set_unary_output_data(in, out);
     int size = in.data_size();
-    out.set_data(
-        allocator::malloc_or_wait(size * out.itemsize()),
-        size,
-        in.strides(),
-        in.flags());
     vvsinhf(out.data<float>(), in.data<float>(), &size);
   } else {
     eval(inputs, out);
@@ -619,12 +515,8 @@ void Square::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   auto& in = inputs[0];
   if (in.dtype() == float32 && in.flags().contiguous) {
+    set_unary_output_data(in, out);
     auto size = in.data_size();
-    out.set_data(
-        allocator::malloc_or_wait(size * out.itemsize()),
-        size,
-        in.strides(),
-        in.flags());
     vDSP_vsq(in.data<float>(), 1, out.data<float>(), 1, size);
   } else {
     unary(in, out, [](auto x) { return x * x; });
@@ -635,12 +527,8 @@ void Sqrt::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   auto& in = inputs[0];
   if (in.dtype() == float32 && in.flags().contiguous) {
+    set_unary_output_data(in, out);
     int size = in.data_size();
-    out.set_data(
-        allocator::malloc_or_wait(size * out.itemsize()),
-        size,
-        in.strides(),
-        in.flags());
     if (recip_) {
       vvrsqrtf(out.data<float>(), in.data<float>(), &size);
     } else {
@@ -695,12 +583,8 @@ void Tan::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   const auto& in = inputs[0];
   if (out.dtype() == float32 && in.flags().contiguous) {
+    set_unary_output_data(in, out);
     int size = in.data_size();
-    out.set_data(
-        allocator::malloc_or_wait(size * out.itemsize()),
-        size,
-        in.strides(),
-        in.flags());
     vvtanf(out.data<float>(), in.data<float>(), &size);
   } else {
     eval(inputs, out);
@@ -711,12 +595,8 @@ void Tanh::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   const auto& in = inputs[0];
   if (out.dtype() == float32 && in.flags().contiguous) {
+    set_unary_output_data(in, out);
     int size = in.data_size();
-    out.set_data(
-        allocator::malloc_or_wait(size * out.itemsize()),
-        size,
-        in.strides(),
-        in.flags());
     vvtanhf(out.data<float>(), in.data<float>(), &size);
   } else {
     eval(inputs, out);

mlx/backend/common/CMakeLists.txt

@@ -16,4 +16,5 @@ target_sources(
   ${CMAKE_CURRENT_SOURCE_DIR}/threefry.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/indexing.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/load.cpp
+  ${CMAKE_CURRENT_SOURCE_DIR}/qrf.cpp
 )

mlx/backend/common/binary.cpp

@@ -233,14 +233,33 @@ void Maximum::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 2);
   auto& a = inputs[0];
   auto& b = inputs[1];
-  binary(a, b, out, [](auto x, auto y) { return (x > y) ? x : y; });
+
+  if (is_floating_point(out.dtype())) {
+    binary(a, b, out, [](auto x, auto y) {
+      if (std::isnan(x)) {
+        return x;
+      }
+      return (x > y) ? x : y;
+    });
+  } else {
+    binary(a, b, out, [](auto x, auto y) { return (x > y) ? x : y; });
+  }
 }
 
 void Minimum::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 2);
   auto& a = inputs[0];
   auto& b = inputs[1];
-  binary(a, b, out, [](auto x, auto y) { return (x < y) ? x : y; });
+  if (is_floating_point(out.dtype())) {
+    binary(a, b, out, [](auto x, auto y) {
+      if (std::isnan(x)) {
+        return x;
+      }
+      return (x < y) ? x : y;
+    });
+  } else {
+    binary(a, b, out, [](auto x, auto y) { return (x < y) ? x : y; });
+  }
 }
 
 void Multiply::eval(const std::vector<array>& inputs, array& out) {

mlx/backend/common/binary.h

@@ -1,7 +1,6 @@
 // Copyright © 2023 Apple Inc.
 
 #pragma once
-
 #include "mlx/allocator.h"
 #include "mlx/array.h"
 #include "mlx/backend/common/utils.h"
@@ -40,29 +39,83 @@ void set_binary_op_output_data(
     const array& a,
     const array& b,
     array& out,
-    BinaryOpType bopt) {
+    BinaryOpType bopt,
+    bool donate_with_move = false) {
   switch (bopt) {
     case ScalarScalar:
       out.set_data(
           allocator::malloc_or_wait(out.itemsize()), 1, a.strides(), a.flags());
       break;
     case ScalarVector:
-      out.set_data(
-          allocator::malloc_or_wait(b.data_size() * out.itemsize()),
-          b.data_size(),
-          b.strides(),
-          b.flags());
+      if (b.is_donatable() && b.itemsize() == out.itemsize()) {
+        if (donate_with_move) {
+          out.move_shared_buffer(b);
+        } else {
+          out.copy_shared_buffer(b);
+        }
+      } else {
+        out.set_data(
+            allocator::malloc_or_wait(b.data_size() * out.itemsize()),
+            b.data_size(),
+            b.strides(),
+            b.flags());
+      }
       break;
     case VectorScalar:
+      if (a.is_donatable() && a.itemsize() == out.itemsize()) {
+        if (donate_with_move) {
+          out.move_shared_buffer(a);
+        } else {
+          out.copy_shared_buffer(a);
+        }
+      } else {
+        out.set_data(
+            allocator::malloc_or_wait(a.data_size() * out.itemsize()),
+            a.data_size(),
+            a.strides(),
+            a.flags());
+      }
+      break;
     case VectorVector:
-      out.set_data(
-          allocator::malloc_or_wait(a.data_size() * out.itemsize()),
-          a.data_size(),
-          a.strides(),
-          a.flags());
+      if (a.is_donatable() && a.itemsize() == out.itemsize()) {
+        if (donate_with_move) {
+          out.move_shared_buffer(a);
+        } else {
+          out.copy_shared_buffer(a);
+        }
+      } else if (b.is_donatable() && b.itemsize() == out.itemsize()) {
+        if (donate_with_move) {
+          out.move_shared_buffer(b);
+        } else {
+          out.copy_shared_buffer(b);
+        }
+      } else {
+        out.set_data(
+            allocator::malloc_or_wait(a.data_size() * out.itemsize()),
+            a.data_size(),
+            a.strides(),
+            a.flags());
+      }
       break;
     case General:
-      out.set_data(allocator::malloc_or_wait(out.nbytes()));
+      if (a.is_donatable() && a.flags().row_contiguous &&
+          a.itemsize() == out.itemsize() && a.size() == out.size()) {
+        if (donate_with_move) {
+          out.move_shared_buffer(a);
+        } else {
+          out.copy_shared_buffer(a);
+        }
+      } else if (
+          b.is_donatable() && b.flags().row_contiguous &&
+          b.itemsize() == out.itemsize() && b.size() == out.size()) {
+        if (donate_with_move) {
+          out.move_shared_buffer(b);
+        } else {
+          out.copy_shared_buffer(b);
+        }
+      } else {
+        out.set_data(allocator::malloc_or_wait(out.nbytes()));
+      }
       break;
   }
 }

mlx/backend/common/copy.cpp

@@ -289,11 +289,16 @@ void copy(const array& src, array& dst, CopyType ctype) {
   // Allocate the output
   switch (ctype) {
     case CopyType::Vector:
-      dst.set_data(
-          allocator::malloc_or_wait(src.data_size() * dst.itemsize()),
-          src.data_size(),
-          src.strides(),
-          src.flags());
+      if (src.is_donatable() && src.itemsize() == dst.itemsize()) {
+        dst.copy_shared_buffer(src);
+      } else {
+        auto size = src.data_size();
+        dst.set_data(
+            allocator::malloc_or_wait(size * dst.itemsize()),
+            size,
+            src.strides(),
+            src.flags());
+      }
       break;
     case CopyType::Scalar:
     case CopyType::General:

mlx/backend/common/default_primitives.cpp

@@ -1,4 +1,4 @@
-// Copyright © 2023 Apple Inc.
+// Copyright © 2023-2024 Apple Inc.
 
 #ifdef ACCELERATE_NEW_LAPACK
 #include <vecLib/cblas_new.h>
@@ -6,6 +6,8 @@
 #include <cblas.h>
 #endif
 
+#include <cstring>
+
 #include "mlx/array.h"
 #include "mlx/backend/common/copy.h"
 #include "mlx/backend/common/utils.h"
@@ -45,6 +47,8 @@ DEFAULT(Convolution)
 DEFAULT(Copy)
 DEFAULT(Cos)
 DEFAULT(Cosh)
+DEFAULT_MULTI(CustomVJP)
+DEFAULT_MULTI(Depends)
 DEFAULT(Divide)
 DEFAULT(Remainder)
 DEFAULT(Equal)
@@ -97,6 +101,7 @@ DEFAULT(Tan)
 DEFAULT(Tanh)
 DEFAULT(Transpose)
 DEFAULT_MULTI(DivMod)
+DEFAULT_MULTI(QRF)
 
 namespace {
 
@@ -127,6 +132,11 @@ inline void matmul_common_general(
   size_t N = b.shape(-1);
   size_t K = a.shape(-1);
 
+  if (K == 0) {
+    std::memset(static_cast<void*>(out.data<float>()), 0, out.nbytes());
+    return;
+  }
+
   for (int i = 0; i < (a.size() / (M * K)); ++i) {
     cblas_sgemm(
         CblasRowMajor,

mlx/backend/common/primitives.cpp

@@ -232,22 +232,38 @@ void Cosh::eval(const std::vector<array>& inputs, array& out) {
   }
 }
 
+void CustomVJP::eval(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  assert(inputs.size() > outputs.size());
+  for (int i = 0, j = inputs.size() - outputs.size(); i < outputs.size();
+       i++, j++) {
+    outputs[i].copy_shared_buffer(inputs[j]);
+  }
+}
+
+void Depends::eval(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  assert(inputs.size() > outputs.size());
+  for (int i = 0; i < outputs.size(); i++) {
+    outputs[i].copy_shared_buffer(inputs[i]);
+  }
+}
+
 void Erf::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   const auto& in = inputs[0];
   switch (out.dtype()) {
     case float32:
-      out.set_data(allocator::malloc_or_wait(out.nbytes()));
       unary_op<float>(in, out, [](auto x) { return std::erf(x); });
       break;
     case float16:
-      out.set_data(allocator::malloc_or_wait(out.nbytes()));
       unary_op<float16_t>(in, out, [](auto x) {
         return static_cast<float16_t>(std::erf(static_cast<float>(x)));
       });
       break;
     case bfloat16:
-      out.set_data(allocator::malloc_or_wait(out.nbytes()));
       unary_op<bfloat16_t>(in, out, [](auto x) {
         return static_cast<bfloat16_t>(std::erf(static_cast<float>(x)));
       });
@@ -264,17 +280,14 @@ void ErfInv::eval(const std::vector<array>& inputs, array& out) {
   const auto& in = inputs[0];
   switch (out.dtype()) {
     case float32:
-      out.set_data(allocator::malloc_or_wait(out.nbytes()));
       unary_op<float>(in, out, [](auto x) { return erfinv(x); });
       break;
     case float16:
-      out.set_data(allocator::malloc_or_wait(out.nbytes()));
       unary_op<float16_t>(in, out, [](auto x) {
         return static_cast<float16_t>(erfinv(static_cast<float>(x)));
       });
       break;
     case bfloat16:
-      out.set_data(allocator::malloc_or_wait(out.nbytes()));
       unary_op<bfloat16_t>(in, out, [](auto x) {
         return static_cast<bfloat16_t>(erfinv(static_cast<float>(x)));
       });

mlx/backend/common/qrf.cpp

@@ -0,0 +1,153 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include "mlx/allocator.h"
+#include "mlx/backend/common/copy.h"
+#include "mlx/primitives.h"
+
+#ifdef ACCELERATE_NEW_LAPACK
+#include <vecLib/lapack.h>
+#else
+#include <lapack.h>
+#endif
+
+namespace mlx::core {
+
+template <typename T>
+struct lpack;
+
+template <>
+struct lpack<float> {
+  static void xgeqrf(
+      const int* m,
+      const int* n,
+      float* a,
+      const int* lda,
+      float* tau,
+      float* work,
+      const int* lwork,
+      int* info) {
+    sgeqrf_(m, n, a, lda, tau, work, lwork, info);
+  }
+  static void xorgqr(
+      const int* m,
+      const int* n,
+      const int* k,
+      float* a,
+      const int* lda,
+      const float* tau,
+      float* work,
+      const int* lwork,
+      int* info) {
+    sorgqr_(m, n, k, a, lda, tau, work, lwork, info);
+  }
+};
+
+template <typename T>
+void qrf_impl(const array& a, array& q, array& r) {
+  const int M = a.shape(-2);
+  const int N = a.shape(-1);
+  const int lda = std::max(M, N);
+  size_t num_matrices = a.size() / (M * N);
+  int num_reflectors = std::min(M, N);
+  auto tau =
+      allocator::malloc_or_wait(sizeof(T) * num_matrices * num_reflectors);
+
+  // Copy A to inplace input and make it col-contiguous
+  array in(a.shape(), float32, nullptr, {});
+  auto flags = in.flags();
+
+  // Copy the input to be column contiguous
+  flags.col_contiguous = num_matrices == 1;
+  flags.row_contiguous = false;
+  std::vector<size_t> strides = in.strides();
+  strides[in.ndim() - 2] = 1;
+  strides[in.ndim() - 1] = M;
+  in.set_data(
+      allocator::malloc_or_wait(in.nbytes()), in.nbytes(), strides, flags);
+  copy_inplace(a, in, CopyType::GeneralGeneral);
+
+  T optimal_work;
+  int lwork = -1;
+  int info;
+
+  // Compute workspace size
+  lpack<T>::xgeqrf(
+      &M, &N, nullptr, &lda, nullptr, &optimal_work, &lwork, &info);
+
+  // Update workspace size
+  lwork = optimal_work;
+  auto work = allocator::malloc_or_wait(sizeof(T) * lwork);
+
+  // Loop over matrices
+  for (int i = 0; i < num_matrices; ++i) {
+    // Solve
+    lpack<T>::xgeqrf(
+        &M,
+        &N,
+        in.data<float>() + M * N * i,
+        &lda,
+        static_cast<T*>(tau.raw_ptr()) + num_reflectors * i,
+        static_cast<T*>(work.raw_ptr()),
+        &lwork,
+        &info);
+  }
+  allocator::free(work);
+
+  r.set_data(allocator::malloc_or_wait(r.nbytes()));
+  copy_inplace(in, r, CopyType::General);
+
+  for (int i = 0; i < num_matrices; ++i) {
+    // Zero lower triangle
+    for (int j = 0; j < r.shape(-2); ++j) {
+      for (int k = 0; k < j; ++k) {
+        r.data<T>()[i * N * M + j * N + k] = 0;
+      }
+    }
+  }
+
+  // Get work size
+  lwork = -1;
+  lpack<T>::xorgqr(
+      &M,
+      &N,
+      &num_reflectors,
+      nullptr,
+      &lda,
+      nullptr,
+      &optimal_work,
+      &lwork,
+      &info);
+  lwork = optimal_work;
+  work = allocator::malloc_or_wait(sizeof(T) * lwork);
+
+  // Loop over matrices
+  for (int i = 0; i < num_matrices; ++i) {
+    // Compute Q
+    lpack<T>::xorgqr(
+        &M,
+        &N,
+        &num_reflectors,
+        in.data<float>() + M * N * i,
+        &lda,
+        static_cast<T*>(tau.raw_ptr()) + num_reflectors * i,
+        static_cast<T*>(work.raw_ptr()),
+        &lwork,
+        &info);
+  }
+
+  q.set_data(allocator::malloc_or_wait(q.nbytes()));
+  copy_inplace(in, q, CopyType::General);
+
+  // Cleanup
+  allocator::free(work);
+  allocator::free(tau);
+}
+
+void QRF::eval(const std::vector<array>& inputs, std::vector<array>& outputs) {
+  if (!(inputs[0].dtype() == float32)) {
+    throw std::runtime_error("[QRF::eval] only supports float32.");
+  }
+  qrf_impl<float>(inputs[0], outputs[0], outputs[1]);
+}
+
+} // namespace mlx::core

mlx/backend/common/quantized.cpp

@@ -1,7 +1,6 @@
 // Copyright © 2023 Apple Inc.
 
 #include <cassert>
-#include <iostream>
 
 #include "mlx/backend/metal/copy.h"
 #include "mlx/primitives.h"
@@ -119,6 +118,12 @@ void _qmm_dispatch_typed(
   switch (bits) {
     case 2: {
       switch (group_size) {
+        case 32:
+          if (transposed_w) {
+            return _qmm_t<T, 2, 32>(result, x, w, scales, biases, M, N, K);
+          } else {
+            return _qmm<T, 2, 32>(result, x, w, scales, biases, M, N, K);
+          }
         case 64:
           if (transposed_w) {
             return _qmm_t<T, 2, 64>(result, x, w, scales, biases, M, N, K);
@@ -135,6 +140,12 @@ void _qmm_dispatch_typed(
     }
     case 4: {
       switch (group_size) {
+        case 32:
+          if (transposed_w) {
+            return _qmm_t<T, 4, 32>(result, x, w, scales, biases, M, N, K);
+          } else {
+            return _qmm<T, 4, 32>(result, x, w, scales, biases, M, N, K);
+          }
         case 64:
           if (transposed_w) {
             return _qmm_t<T, 4, 64>(result, x, w, scales, biases, M, N, K);
@@ -151,6 +162,12 @@ void _qmm_dispatch_typed(
     }
     case 8: {
       switch (group_size) {
+        case 32:
+          if (transposed_w) {
+            return _qmm_t<T, 8, 32>(result, x, w, scales, biases, M, N, K);
+          } else {
+            return _qmm<T, 8, 32>(result, x, w, scales, biases, M, N, K);
+          }
         case 64:
           if (transposed_w) {
             return _qmm_t<T, 8, 64>(result, x, w, scales, biases, M, N, K);

mlx/backend/common/unary.h

@@ -64,15 +64,24 @@ struct RoundOp {
   }
 };
 
+void set_unary_output_data(const array& in, array& out) {
+  if (in.is_donatable() && in.itemsize() == out.itemsize()) {
+    out.copy_shared_buffer(in);
+  } else {
+    auto size = in.data_size();
+    out.set_data(
+        allocator::malloc_or_wait(size * out.itemsize()),
+        size,
+        in.strides(),
+        in.flags());
+  }
+}
+
 template <typename T, typename Op>
 void unary_op(const array& a, array& out, Op op) {
   const T* a_ptr = a.data<T>();
   if (a.flags().contiguous) {
-    out.set_data(
-        allocator::malloc_or_wait(a.data_size() * out.itemsize()),
-        a.data_size(),
-        a.strides(),
-        a.flags());
+    set_unary_output_data(a, out);
     T* dst = out.data<T>();
     for (size_t i = 0; i < a.data_size(); ++i) {
       dst[i] = op(a_ptr[i]);

mlx/backend/metal/conv.cpp

@@ -2,7 +2,6 @@
 
 #include <algorithm>
 #include <cassert>
-#include <iostream>
 #include <numeric>
 #include <sstream>
 

mlx/backend/metal/copy.cpp

@@ -12,11 +12,15 @@ namespace mlx::core {
 
 void copy_gpu(const array& in, array& out, CopyType ctype, const Stream& s) {
   if (ctype == CopyType::Vector) {
-    out.set_data(
-        allocator::malloc_or_wait(in.data_size() * out.itemsize()),
-        in.data_size(),
-        in.strides(),
-        in.flags());
+    if (in.is_donatable() && in.itemsize() == out.itemsize()) {
+      out.move_shared_buffer(in);
+    } else {
+      out.set_data(
+          allocator::malloc_or_wait(in.data_size() * out.itemsize()),
+          in.data_size(),
+          in.strides(),
+          in.flags());
+    }
   } else {
     out.set_data(allocator::malloc_or_wait(out.nbytes()));
   }
@@ -67,7 +71,8 @@ void copy_gpu_inplace(
   auto kernel = d.get_kernel(kname.str());
   auto compute_encoder = d.get_command_encoder(s.index);
   compute_encoder->setComputePipelineState(kernel);
-  set_array_buffer(compute_encoder, in, 0);
+  bool donate_in = in.data_shared_ptr() == nullptr;
+  set_array_buffer(compute_encoder, donate_in ? out : in, 0);
   set_array_buffer(compute_encoder, out, 1);
 
   if (ctype == CopyType::General || ctype == CopyType::GeneralGeneral) {

mlx/backend/metal/device.cpp

@@ -1,4 +1,4 @@
-// Copyright © 2023 Apple Inc.
+// Copyright © 2023-24 Apple Inc.
 
 #include <dlfcn.h>
 #include <cstdlib>
@@ -242,37 +242,127 @@ void Device::register_library(
   }
 }
 
-MTL::ComputePipelineState* Device::get_kernel(
-    const std::string& name,
-    const std::string& lib_name /* = "mlx" */) {
-  auto pool = new_scoped_memory_pool();
-  // Look for cached kernel
-  if (auto it = kernel_map_.find(name); it != kernel_map_.end()) {
-    return it->second;
-  }
-
-  // Prepare new kernel
-
+MTL::Library* Device::get_library_cache_(const std::string& lib_name) {
   // Search for cached metal lib
   MTL::Library* mtl_lib;
-  if (auto it = library_map_.find(name); it != library_map_.end()) {
+  if (auto it = library_map_.find(lib_name); it != library_map_.end()) {
     mtl_lib = it->second;
   } else { // Look for metallib alongside library
     register_library(lib_name);
     mtl_lib = library_map_[lib_name];
   }
 
+  return mtl_lib;
+}
+
+MTL::Library* Device::get_library_(const std::string& source_string) {
+  auto pool = new_scoped_memory_pool();
+
+  auto ns_code =
+      NS::String::string(source_string.c_str(), NS::ASCIIStringEncoding);
+
+  NS::Error* error = nullptr;
+  auto mtl_lib = device_->newLibrary(ns_code, nullptr, &error);
+
+  // Throw error if unable to compile library
+  if (!mtl_lib) {
+    std::ostringstream msg;
+    msg << "[metal::Device] Unable to load build metal library from source"
+        << "\n";
+    if (error) {
+      msg << error->localizedDescription()->utf8String() << "\n";
+    }
+    throw std::runtime_error(msg.str());
+  }
+
+  return mtl_lib;
+}
+
+MTL::Library* Device::get_library_(const MTL::StitchedLibraryDescriptor* desc) {
+  auto pool = new_scoped_memory_pool();
+
+  NS::Error* error = nullptr;
+  auto mtl_lib = device_->newLibrary(desc, &error);
+
+  // Throw error if unable to compile library
+  if (!mtl_lib) {
+    std::ostringstream msg;
+    msg << "[metal::Device] Unable to load build stitched metal library"
+        << "\n";
+    if (error) {
+      msg << error->localizedDescription()->utf8String() << "\n";
+    }
+    throw std::runtime_error(msg.str());
+  }
+
+  return mtl_lib;
+}
+
+MTL::Function* Device::get_function_(
+    const std::string& name,
+    MTL::Library* mtl_lib) {
   // Pull kernel from library
   auto ns_name = NS::String::string(name.c_str(), NS::ASCIIStringEncoding);
   auto mtl_function = mtl_lib->newFunction(ns_name);
 
+  return mtl_function;
+}
+
+MTL::Function* Device::get_function_(
+    const std::string& name,
+    const std::string& specialized_name,
+    const MTLFCList& func_consts,
+    MTL::Library* mtl_lib) {
+  if (func_consts.empty() && (specialized_name == name)) {
+    return get_function_(name, mtl_lib);
+  }
+
+  // Prepare function constants
+  auto mtl_func_consts = MTL::FunctionConstantValues::alloc()->init();
+
+  for (auto [value, type, index] : func_consts) {
+    mtl_func_consts->setConstantValue(value, type, index);
+  }
+
+  // Prepare function desc
+  auto desc = MTL::FunctionDescriptor::functionDescriptor();
+  desc->setName(NS::String::string(name.c_str(), NS::ASCIIStringEncoding));
+  desc->setSpecializedName(
+      NS::String::string(specialized_name.c_str(), NS::ASCIIStringEncoding));
+  desc->setConstantValues(mtl_func_consts);
+
+  // Pull kernel from library
+  NS::Error* error = nullptr;
+  auto mtl_function = mtl_lib->newFunction(desc, &error);
+
+  // Throw error if unable to build metal function
+  if (!mtl_function) {
+    std::ostringstream msg;
+    msg << "[metal::Device] Unable to load function " << name << "\n";
+    if (error) {
+      msg << error->localizedDescription()->utf8String() << "\n";
+    }
+    throw std::runtime_error(msg.str());
+  }
+
+  mtl_func_consts->release();
+  desc->release();
+
+  return mtl_function;
+}
+
+MTL::ComputePipelineState* Device::get_kernel_(
+    const std::string& name,
+    const MTL::Function* mtl_function) {
   // Compile kernel to compute pipeline
   NS::Error* error = nullptr;
   MTL::ComputePipelineState* kernel;
+
   if (mtl_function) {
     kernel = device_->newComputePipelineState(mtl_function, &error);
-    mtl_function->release();
   }
+
+  // Throw error if unable to compile metal function
   if (!mtl_function || !kernel) {
     std::ostringstream msg;
     msg << "[metal::Device] Unable to load kernel " << name << "\n";
@@ -282,11 +372,170 @@ MTL::ComputePipelineState* Device::get_kernel(
     throw std::runtime_error(msg.str());
   }
 
-  // Add kernel to cache
-  kernel_map_.insert({name, kernel});
   return kernel;
 }
 
+MTL::ComputePipelineState* Device::get_kernel_(
+    const std::string& name,
+    const MTL::Function* mtl_function,
+    const MTL::LinkedFunctions* linked_functions) {
+  // Check inputs
+  if (!linked_functions) {
+    return get_kernel_(name, mtl_function);
+  }
+
+  if (!mtl_function) {
+    std::ostringstream msg;
+    msg << "[metal::Device] Unable to load kernel " << name << "\n";
+    throw std::runtime_error(msg.str());
+  }
+
+  // Prepare compute pipeline state descriptor
+  auto desc = MTL::ComputePipelineDescriptor::alloc()->init();
+  desc->setComputeFunction(mtl_function);
+  desc->setLinkedFunctions(linked_functions);
+
+  // Compile kernel to compute pipeline
+  NS::Error* error = nullptr;
+  auto kernel = device_->newComputePipelineState(
+      desc, MTL::PipelineOptionNone, nullptr, &error);
+
+  // Throw error if unable to compile metal function
+  if (!kernel) {
+    std::ostringstream msg;
+    msg << "[metal::Device] Unable to load kernel " << name << "\n";
+    if (error) {
+      msg << error->localizedDescription()->utf8String() << "\n";
+    }
+    throw std::runtime_error(msg.str());
+  }
+
+  return kernel;
+}
+
+MTL::Library* Device::get_library(
+    const std::string& name,
+    const std::string& source,
+    bool cache /* = true */) {
+  if (cache) {
+    if (auto it = library_map_.find(name); it != library_map_.end()) {
+      return it->second;
+    }
+  }
+
+  auto mtl_lib = get_library_(source);
+
+  if (cache) {
+    library_map_.insert({name, mtl_lib});
+  }
+
+  return mtl_lib;
+}
+
+MTL::Library* Device::get_library(
+    const std::string& name,
+    const MTL::StitchedLibraryDescriptor* desc,
+    bool cache /* = true */) {
+  if (cache) {
+    if (auto it = library_map_.find(name); it != library_map_.end()) {
+      return it->second;
+    }
+  }
+
+  auto mtl_lib = get_library_(desc);
+
+  if (cache) {
+    library_map_.insert({name, mtl_lib});
+  }
+
+  return mtl_lib;
+}
+
+MTL::Function* Device::get_function(
+    const std::string& base_name,
+    MTL::Library* mtl_lib,
+    const std::string& specialized_name /* = "" */,
+    const MTLFCList& func_consts /* = {} */) {
+  return get_function_(base_name, specialized_name, func_consts, mtl_lib);
+}
+
+MTL::Function* Device::get_function(
+    const std::string& base_name,
+    const std::string& lib_name /* = "mlx" */,
+    const std::string& specialized_name /*  = "" */,
+    const MTLFCList& func_consts /* = {} */) {
+  // Search for cached metal lib
+  MTL::Library* mtl_lib = get_library_cache_(lib_name);
+
+  return get_function(base_name, mtl_lib, specialized_name, func_consts);
+}
+
+MTL::LinkedFunctions* Device::get_linked_functions_(
+    const std::vector<MTL::Function*>& funcs) {
+  if (funcs.empty()) {
+    return nullptr;
+  }
+
+  auto lfuncs = MTL::LinkedFunctions::linkedFunctions();
+
+  std::vector<NS::Object*> objs(funcs.size());
+  for (int i = 0; i < funcs.size(); i++) {
+    objs[i] = funcs[i];
+  }
+
+  NS::Array* funcs_arr = NS::Array::array(objs.data(), funcs.size());
+
+  lfuncs->setPrivateFunctions(funcs_arr);
+
+  return lfuncs;
+}
+
+MTL::ComputePipelineState* Device::get_kernel(
+    const std::string& base_name,
+    MTL::Library* mtl_lib,
+    const std::string& hash_name /* = "" */,
+    const MTLFCList& func_consts /* = {} */,
+    const std::vector<MTL::Function*>& linked_functions /* = {} */) {
+  auto pool = new_scoped_memory_pool();
+
+  // Look for cached kernel
+  const auto& kname = hash_name.empty() ? base_name : hash_name;
+  if (auto it = kernel_map_.find(kname); it != kernel_map_.end()) {
+    return it->second;
+  }
+
+  // Pull kernel from library
+  auto mtl_function = get_function_(base_name, kname, func_consts, mtl_lib);
+
+  // Compile kernel to compute pipeline
+  auto mtl_linked_funcs = get_linked_functions_(linked_functions);
+  auto kernel = get_kernel_(kname, mtl_function, mtl_linked_funcs);
+  mtl_function->release();
+  mtl_linked_funcs->release();
+
+  // Add kernel to cache
+  kernel_map_.insert({kname, kernel});
+  return kernel;
+}
+
+MTL::ComputePipelineState* Device::get_kernel(
+    const std::string& base_name,
+    const std::string& lib_name /* = "mlx" */,
+    const std::string& hash_name /*  = "" */,
+    const MTLFCList& func_consts /*  = {} */,
+    const std::vector<MTL::Function*>& linked_functions /*  = {} */) {
+  // Look for cached kernel
+  const auto& kname = hash_name.size() == 0 ? base_name : hash_name;
+  if (auto it = kernel_map_.find(kname); it != kernel_map_.end()) {
+    return it->second;
+  }
+
+  // Search for cached metal lib
+  MTL::Library* mtl_lib = get_library_cache_(lib_name);
+
+  return get_kernel(base_name, mtl_lib, kname, func_consts, linked_functions);
+}
+
 Device& device(mlx::core::Device) {
   static Device metal_device;
   return metal_device;

mlx/backend/metal/device.h

@@ -1,4 +1,4 @@
-// Copyright © 2023 Apple Inc.
+// Copyright © 2023-24 Apple Inc.
 
 #pragma once
 
@@ -31,6 +31,9 @@ inline std::string get_colocated_mtllib_path(const std::string& lib_name) {
   return mtllib_path;
 }
 
+using MTLFCList =
+    std::vector<std::tuple<const void*, MTL::DataType, NS::UInteger>>;
+
 class Device {
  public:
   Device();
@@ -59,14 +62,71 @@ class Device {
       const std::function<std::string(const std::string&)>& lib_path_func =
           get_colocated_mtllib_path);
 
-  MTL::ComputePipelineState* get_kernel(
+  MTL::Library* get_library(
       const std::string& name,
-      const std::string& lib_name = "mlx");
+      const std::string& source_string,
+      bool cache = true);
+
+  MTL::Library* get_library(
+      const std::string& name,
+      const MTL::StitchedLibraryDescriptor* desc,
+      bool cache = true);
+
+  MTL::Function* get_function(
+      const std::string& base_name,
+      MTL::Library* mtl_lib,
+      const std::string& specialized_name = "",
+      const MTLFCList& func_consts = {});
+
+  MTL::Function* get_function(
+      const std::string& base_name,
+      const std::string& lib_name = "mlx",
+      const std::string& specialized_name = "",
+      const MTLFCList& func_consts = {});
+
+  MTL::ComputePipelineState* get_kernel(
+      const std::string& base_name,
+      MTL::Library* mtl_lib,
+      const std::string& hash_name = "",
+      const MTLFCList& func_consts = {},
+      const std::vector<MTL::Function*>& linked_functions = {});
+
+  MTL::ComputePipelineState* get_kernel(
+      const std::string& base_name,
+      const std::string& lib_name = "mlx",
+      const std::string& hash_name = "",
+      const MTLFCList& func_consts = {},
+      const std::vector<MTL::Function*>& linked_functions = {});
 
   MTL::ArgumentEncoder* argument_encoder(
       const std::vector<MTL::ArgumentDescriptor*>& arg_descs) const;
 
  private:
+  MTL::Library* get_library_cache_(const std::string& name);
+
+  MTL::Library* get_library_(const std::string& source_string);
+  MTL::Library* get_library_(const MTL::StitchedLibraryDescriptor* desc);
+
+  MTL::Function* get_function_(const std::string& name, MTL::Library* mtl_lib);
+
+  MTL::Function* get_function_(
+      const std::string& name,
+      const std::string& specialized_name,
+      const MTLFCList& func_consts,
+      MTL::Library* mtl_lib);
+
+  MTL::LinkedFunctions* get_linked_functions_(
+      const std::vector<MTL::Function*>& funcs);
+
+  MTL::ComputePipelineState* get_kernel_(
+      const std::string& name,
+      const MTL::Function* mtl_function);
+
+  MTL::ComputePipelineState* get_kernel_(
+      const std::string& name,
+      const MTL::Function* mtl_function,
+      const MTL::LinkedFunctions* linked_functions);
+
   MTL::Device* device_;
   std::unordered_map<int32_t, MTL::CommandQueue*> queue_map_;
   std::unordered_map<int32_t, std::pair<int, MTL::CommandBuffer*>> buffer_map_;

mlx/backend/metal/kernels/arg_reduce.metal

@@ -63,18 +63,6 @@ struct ArgMax {
   }
 };
 
-bool simd_shuffle_down(bool data, uint16_t delta) {
-  return simd_shuffle_down(static_cast<uint32_t>(data), delta);
-}
-
-uint64_t simd_shuffle_down(uint64_t data, uint16_t delta) {
-  return as_type<uint64_t>(simd_shuffle_down(as_type<uint2>(data), delta));
-}
-
-int64_t simd_shuffle_down(int64_t data, uint16_t delta) {
-  return as_type<int64_t>(simd_shuffle_down(as_type<uint2>(data), delta));
-}
-
 template <typename U>
 IndexValPair<U> simd_shuffle_down(IndexValPair<U> data, uint16_t delta) {
   return IndexValPair<U>(

mlx/backend/metal/kernels/binary.metal

@@ -58,6 +58,9 @@ struct LessEqual {
 struct LogAddExp {
   template <typename T>
   T operator()(T x, T y) {
+    if (metal::isnan(x) || metal::isnan(y)) {
+      return metal::numeric_limits<T>::quiet_NaN();
+    }
     constexpr T inf = metal::numeric_limits<T>::infinity();
     T maxval = metal::max(x, y);
     T minval = metal::min(x, y);
@@ -67,20 +70,48 @@ struct LogAddExp {
 };
 
 struct Maximum {
-  template <typename T> T operator()(T x, T y) { return metal::max(x, y); }
+  template <typename T>
+  metal::enable_if_t<metal::is_integral_v<T>, T> operator()(T x, T y) {
+    return metal::max(x, y);
+  }
+
+  template <typename T>
+  metal::enable_if_t<!metal::is_integral_v<T>, T> operator()(T x, T y) {
+    if (metal::isnan(x)) {
+      return x;
+    }
+    return x > y ? x : y;
+  }
 
   template <>
   complex64_t operator()(complex64_t x, complex64_t y) {
-    return x >= y ? x : y;
+    if (metal::isnan(x.real) || metal::isnan(x.imag)) {
+      return x;
+    }
+    return x > y ? x : y;
   }
 };
 
 struct Minimum {
-  template <typename T> T operator()(T x, T y) { return metal::min(x, y); }
+  template <typename T>
+  metal::enable_if_t<metal::is_integral_v<T>, T> operator()(T x, T y) {
+    return metal::min(x, y);
+  }
+
+  template <typename T>
+  metal::enable_if_t<!metal::is_integral_v<T>, T> operator()(T x, T y) {
+    if (metal::isnan(x)) {
+      return x;
+    }
+    return x < y ? x : y;
+  }
 
   template <>
   complex64_t operator()(complex64_t x, complex64_t y) {
-    return x <= y ? x : y;
+    if (metal::isnan(x.real) || metal::isnan(x.imag)) {
+      return x;
+    }
+    return x < y ? x : y;
   }
 };
 
@@ -389,4 +420,4 @@ instantiate_binary_all(naneq, bfloat16, bfloat16_t, bool, NaNEqual)
 instantiate_binary_all(naneq, complex64, complex64_t, bool, NaNEqual)
 
 instantiate_binary_all(lor, bool_, bool, bool, LogicalOr)
-instantiate_binary_all(land, bool_, bool, bool, LogicalAnd)
+instantiate_binary_all(land, bool_, bool, bool, LogicalAnd)

mlx/backend/metal/kernels/gemv.metal

@@ -121,8 +121,18 @@ struct GEMVKernel {
       for(int tm = 0; tm < TM; tm++) {
 
         // Load for the row 
-        for(int tn = 0; tn < TN; tn++) {
-          inter[tn] = mat[tm * in_vec_size + bn + tn];
+        if(bn + TN <= in_vec_size) {
+          #pragma clang loop unroll(full)
+          for(int tn = 0; tn < TN; tn++) {
+            inter[tn] = mat[tm * in_vec_size + bn + tn];
+          }
+
+        } else { // Edgecase
+          #pragma clang loop unroll(full)
+          for(int tn = 0; tn < TN; tn++) {
+            int col_idx = (bn + tn) < in_vec_size ? (bn + tn) : (in_vec_size - 1);
+            inter[tn] = mat[tm * in_vec_size + col_idx];
+          }
         }
 
         // Accumulate results

mlx/backend/metal/kernels/quantized.metal

@@ -142,10 +142,11 @@ template <typename T, const int BM, const int BN, const int group_size, const in
   // Adjust positions
   const int out_vec_size_w = out_vec_size / el_per_int;
   const int out_vec_size_g = out_vec_size / group_size;
-  int out_col = (tid.y * BN + simd_gid) * el_per_int;
+  int out_col_start = tid.y * (BN * el_per_int);
+  int out_col = out_col_start + simd_gid * el_per_int;
   w += out_col / el_per_int;
-  scales += out_col / group_size;
-  biases += out_col / group_size;
+  scales += out_col_start / group_size;
+  biases += out_col_start / group_size;
   x += tid.z * in_vec_size;
   y += tid.z * out_vec_size + out_col;
 
@@ -155,26 +156,22 @@ template <typename T, const int BM, const int BN, const int group_size, const in
 
   // Loop over in_vec in blocks of colgroup
   for (int i=0; i<in_vec_size; i+=BM) {
-    int offset = simd_lid + i;
-    bool thread_in_bounds = offset < in_vec_size;
+    int offset_lid = simd_lid + i;
+    int offset_gid = simd_gid + i;
+    bool thread_in_bounds = offset_lid < in_vec_size;
+    bool group_in_bounds = offset_gid < in_vec_size;
 
     // Load the vec to shared memory
     threadgroup_barrier(mem_flags::mem_threadgroup);
     if (simd_gid == 0) {
-      x_block[simd_lid] = (thread_in_bounds) ? x[offset] : 0;
+      x_block[simd_lid] = (thread_in_bounds) ? x[offset_lid] : 0;
     }
 
     // Load the scales and biases to shared memory
     threadgroup_barrier(mem_flags::mem_threadgroup);
-    if (simd_gid == 0) {
-      #pragma clang loop unroll(full)
-      for (int j=0; j<groups_per_block; j++) {
-        scales_block[simd_lid * groups_per_block + j] = scales[(i + simd_lid) * out_vec_size_g + j];
-      }
-      #pragma clang loop unroll(full)
-      for (int j=0; j<groups_per_block; j++) {
-        biases_block[simd_lid * groups_per_block + j] = biases[(i + simd_lid) * out_vec_size_g + j];
-      }
+    if (simd_lid < groups_per_block && group_in_bounds) {
+      scales_block[simd_gid * groups_per_block + simd_lid] = scales[offset_gid * out_vec_size_g + simd_lid];
+      biases_block[simd_gid * groups_per_block + simd_lid] = biases[offset_gid * out_vec_size_g + simd_lid];
     }
     threadgroup_barrier(mem_flags::mem_threadgroup);
 
@@ -184,7 +181,7 @@ template <typename T, const int BM, const int BN, const int group_size, const in
     bias = biases_block[simd_lid * groups_per_block + (simd_gid * el_per_int) / group_size];
 
     // Load the matrix elements
-    w_local = (thread_in_bounds) ? w[offset * out_vec_size_w] : 0;
+    w_local = (thread_in_bounds) ? w[offset_lid * out_vec_size_w] : 0;
 
     // Do all the work.
     #pragma clang loop unroll(full)
@@ -543,6 +540,9 @@ instantiate_qmv_types(128, 8)
 instantiate_qmv_types( 64, 2)
 instantiate_qmv_types( 64, 4)
 instantiate_qmv_types( 64, 8)
+instantiate_qmv_types( 32, 2)
+instantiate_qmv_types( 32, 4)
+instantiate_qmv_types( 32, 8)
 
 #define instantiate_qvm(name, itype, group_size, bits) \
   template [[host_name("qvm_" #name "_gs_" #group_size "_b_" #bits)]] \
@@ -570,6 +570,9 @@ instantiate_qvm_types(128, 8)
 instantiate_qvm_types( 64, 2)
 instantiate_qvm_types( 64, 4)
 instantiate_qvm_types( 64, 8)
+instantiate_qvm_types( 32, 2)
+instantiate_qvm_types( 32, 4)
+instantiate_qvm_types( 32, 8)
 
 #define instantiate_qmm_t(name, itype, group_size, bits, aligned_N) \
   template [[host_name("qmm_t_" #name "_gs_" #group_size "_b_" #bits "_alN_" #aligned_N)]] \
@@ -601,6 +604,9 @@ instantiate_qmm_t_types(128, 8)
 instantiate_qmm_t_types( 64, 2)
 instantiate_qmm_t_types( 64, 4)
 instantiate_qmm_t_types( 64, 8)
+instantiate_qmm_t_types( 32, 2)
+instantiate_qmm_t_types( 32, 4)
+instantiate_qmm_t_types( 32, 8)
 
 #define instantiate_qmm_n(name, itype, group_size, bits) \
   template [[host_name("qmm_n_" #name "_gs_" #group_size "_b_" #bits)]] \
@@ -629,3 +635,6 @@ instantiate_qmm_n_types(128, 8)
 instantiate_qmm_n_types( 64, 2)
 instantiate_qmm_n_types( 64, 4)
 instantiate_qmm_n_types( 64, 8)
+instantiate_qmm_n_types( 32, 2)
+instantiate_qmm_n_types( 32, 4)
+instantiate_qmm_n_types( 32, 8)

mlx/backend/metal/kernels/reduce.metal

@@ -24,11 +24,59 @@ template <typename T, typename Op>
       device otype *out [[buffer(1)]], \
       uint tid [[thread_position_in_grid]]);
 
-
 ///////////////////////////////////////////////////////////////////////////////
 // All reduce
 ///////////////////////////////////////////////////////////////////////////////
 
+template <typename T, typename U, typename Op, int N_READS=REDUCE_N_READS>
+inline U per_thread_all_reduce(
+    const device T *in,
+    const device size_t& in_size,
+    uint gid,
+    uint grid_size) {
+  Op op;
+  U total_val = Op::init;
+
+  if (gid * N_READS < in_size) {
+    in += gid * N_READS;
+
+    int r = 0;
+    for(; r < (int)ceildiv(in_size, grid_size * N_READS) - 1; r++) {
+      U vals[N_READS] = {op.init};
+
+      for(int i = 0; i < N_READS; i++) {
+        vals[i] = static_cast<U>(in[i]);
+      }
+      for(int i = 0; i < N_READS; i++) {
+        total_val = op(vals[i], total_val);
+      }
+
+      in += grid_size * N_READS;
+    }
+
+    // Separate case for the last set as we close the reduction size
+    size_t curr_idx = (gid + r * (size_t)grid_size) * N_READS;
+    if (curr_idx < in_size) {
+      int max_reads = in_size - curr_idx;
+      T vals[N_READS];
+
+      for(int i = 0, idx = 0; i < N_READS; i++, idx++) {
+        idx = idx < max_reads ? idx : max_reads - 1;
+        vals[i] = in[idx];
+      }
+      for(int i = 0; i < N_READS; i++) {
+        U val = i < max_reads ? vals[i] : Op::init;
+        total_val = op(static_cast<U>(val), total_val);
+      }
+    }
+  }
+
+  return total_val;
+}
+
+// NB: This kernel assumes threads_per_threadgroup is at most
+// 1024. This way with a simd_size of 32, we are guaranteed to
+// complete the reduction in two steps of simd-level reductions.
 template <typename T, typename U, typename Op, int N_READS=REDUCE_N_READS>
 [[kernel]] void all_reduce(
     const device T *in [[buffer(0)]],
@@ -40,53 +88,18 @@ template <typename T, typename U, typename Op, int N_READS=REDUCE_N_READS>
     uint simd_per_group [[simdgroups_per_threadgroup]],
     uint simd_lane_id [[thread_index_in_simdgroup]],
     uint simd_group_id [[simdgroup_index_in_threadgroup]]) {
-  // NB: this kernel assumes threads_per_threadgroup is at most
-  // 1024. This way with a simd_size of 32, we are guaranteed to
-  // complete the reduction in two steps of simd-level reductions.
 
   Op op;
   threadgroup U local_vals[simd_size];
-      
-  U total_val = Op::init;
 
-  in += gid * N_READS;
-      
-  int r = 0;
-  for(; r < (int)ceildiv(in_size, grid_size * N_READS) - 1; r++) {
-    U vals[N_READS] = {op.init}; 
-
-    for(int i = 0; i < N_READS; i++) {
-      vals[i] = static_cast<U>(in[i]);
-    }
-    for(int i = 0; i < N_READS; i++) {
-      total_val = op(vals[i], total_val);
-    }
-
-    in += grid_size * N_READS;
-  }
-
-  // Separate case for the last set as we close the reduction size 
-  size_t curr_idx = (gid + r * (size_t)grid_size) * N_READS;
-  if (curr_idx < in_size) {
-    int max_reads = in_size - curr_idx;
-    T vals[N_READS];
-
-    for(int i = 0, idx = 0; i < N_READS; i++, idx++) {
-      idx = idx < max_reads ? idx : max_reads - 1;
-      vals[i] = in[idx];
-    }
-    for(int i = 0; i < N_READS; i++) {
-      U val = i < max_reads ? vals[i] : Op::init; 
-      total_val = op(static_cast<U>(val), total_val);
-    }
-  }
+  U total_val = per_thread_all_reduce<T, U, Op, N_READS>(in, in_size, gid, grid_size);
 
   // Reduction within simd group
   total_val = op.simd_reduce(total_val);
   if (simd_lane_id == 0) {
     local_vals[simd_group_id] = total_val;
   }
-  
+
   // Reduction within thread group
   threadgroup_barrier(mem_flags::mem_threadgroup);
   total_val = lid < simd_per_group ? local_vals[lid] : op.init;
@@ -98,6 +111,46 @@ template <typename T, typename U, typename Op, int N_READS=REDUCE_N_READS>
   }
 }
 
+template <typename T, typename U, typename Op, int N_READS=REDUCE_N_READS>
+[[kernel]] void all_reduce_no_atomics(
+    const device T *in [[buffer(0)]],
+    device U *out [[buffer(1)]],
+    const device size_t& in_size [[buffer(2)]],
+    uint gid [[thread_position_in_grid]],
+    uint lid [[thread_position_in_threadgroup]],
+    uint grid_size [[threads_per_grid]],
+    uint simd_per_group [[simdgroups_per_threadgroup]],
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]],
+    uint thread_group_id [[threadgroup_position_in_grid]]) {
+
+  Op op;
+  threadgroup U local_vals[simd_size];
+
+  U total_val = per_thread_all_reduce<T, U, Op, N_READS>(in, in_size, gid, grid_size);
+
+  // Reduction within simd group (simd_add isn't supported for uint64/int64 types)
+  for (uint16_t lane_offset = simd_size/2; lane_offset > 0; lane_offset /= 2) {
+    total_val = op(total_val, simd_shuffle_down(total_val, lane_offset));
+  }
+  // Write simd group reduction results to local memory
+  if (simd_lane_id == 0) {
+    local_vals[simd_group_id] = total_val;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  // Reduction of simdgroup reduction results within threadgroup.
+  total_val = lid < simd_per_group ? local_vals[lid] : op.init;
+  for (uint16_t lane_offset = simd_size/2; lane_offset > 0; lane_offset /= 2) {
+    total_val = op(total_val, simd_shuffle_down(total_val, lane_offset));
+  }
+
+  // Reduction across threadgroups
+  if (lid == 0) {
+    out[thread_group_id] = total_val;
+  }
+}
+
 #define instantiate_all_reduce(name, itype, otype, op) \
   template [[host_name("all_reduce_" #name)]] \
   [[kernel]] void all_reduce<itype, otype, op>( \
@@ -111,11 +164,80 @@ template <typename T, typename U, typename Op, int N_READS=REDUCE_N_READS>
       uint simd_lane_id [[thread_index_in_simdgroup]], \
       uint simd_group_id [[simdgroup_index_in_threadgroup]]);
 
+#define instantiate_all_reduce_no_atomics(name, itype, otype, op) \
+  template [[host_name("all_reduce_no_atomics_" #name)]] \
+  [[kernel]] void all_reduce_no_atomics<itype, otype, op>( \
+      const device itype *in [[buffer(0)]], \
+      device otype *out [[buffer(1)]], \
+      const device size_t& in_size [[buffer(2)]], \
+      uint gid [[thread_position_in_grid]], \
+      uint lid [[thread_position_in_threadgroup]], \
+      uint grid_size [[threads_per_grid]], \
+      uint simd_per_group [[simdgroups_per_threadgroup]], \
+      uint simd_lane_id [[thread_index_in_simdgroup]], \
+      uint simd_group_id [[simdgroup_index_in_threadgroup]], \
+      uint thread_group_id [[threadgroup_position_in_grid]]);
 
 ///////////////////////////////////////////////////////////////////////////////
 // Row atomics
 ///////////////////////////////////////////////////////////////////////////////
 
+template <typename T, typename U, typename Op, int N_READS=REDUCE_N_READS>
+inline U per_thread_row_reduce(
+    const device T *in,
+    const constant size_t& reduction_size,
+    const constant size_t& out_size,
+    const constant int* shape,
+    const constant size_t* strides,
+    const constant int& ndim,
+    uint lsize_x,
+    uint lid_x,
+    uint2 tid) {
+
+  Op op;
+
+  // Each threadgroup handles 1 reduction
+  // TODO: Specializing elem_to_loc would be slightly faster
+  int idx = tid.y * out_size + tid.x;
+  int extra_offset = elem_to_loc(idx, shape, strides, ndim);
+  in += extra_offset + lid_x * N_READS;
+  
+  // The reduction is accumulated here
+  U total_val = Op::init;
+
+  // Loop over the reduction size within thread group
+  int r = 0;
+  for (; r < (int)ceildiv(reduction_size, N_READS*lsize_x) - 1; r++) {
+    T vals[N_READS]; 
+    for(int i = 0; i < N_READS; i++) {
+      vals[i] = in[i];
+    }
+    for(int i = 0; i < N_READS; i++) {
+      total_val = op(static_cast<U>(vals[i]), total_val);
+    }
+
+    in += lsize_x * N_READS;
+  }
+
+  // Separate case for the last set as we close the reduction size   
+  size_t reduction_index = (lid_x + (size_t)lsize_x * r) * N_READS;
+  if(reduction_index < reduction_size) {
+    int max_reads = reduction_size - reduction_index;
+
+    T vals[N_READS]; 
+    for(int i = 0; i < N_READS; i++) {
+      int idx = min(i, max_reads - 1);
+      vals[i] = static_cast<U>(in[idx]);
+    }
+    for(int i = 0; i < N_READS; i++) {
+      T val = i < max_reads ? vals[i] : Op::init;
+      total_val = op(static_cast<U>(val), total_val);
+    }
+  }
+
+  return total_val;
+}
+
 template <typename T, typename U, typename Op, int N_READS=REDUCE_N_READS>
 [[kernel]] void row_reduce_general(
     const device T *in [[buffer(0)]],
@@ -133,46 +255,9 @@ template <typename T, typename U, typename Op, int N_READS=REDUCE_N_READS>
     uint simd_group_id [[simdgroup_index_in_threadgroup]]) {
 
   Op op;
-
-  // Each threadgroup handles 1 reduction
-  // TODO: Specializing elem_to_loc would be slightly faster
-  int idx = tid.y * out_size + tid.x;
-  int extra_offset = elem_to_loc(idx, shape, strides, ndim);
-  in += extra_offset + lid.x * N_READS;
-  
-  // The reduction is accumulated here
-  U total_val = Op::init;
   threadgroup U local_vals[simd_size];
 
-  // Loop over the reduction size within thread group
-  int r = 0;
-  for (; r < (int)ceildiv(reduction_size, N_READS*lsize.x) - 1; r++) {
-    T vals[N_READS]; 
-    for(int i = 0; i < N_READS; i++) {
-      vals[i] = in[i];
-    }
-    for(int i = 0; i < N_READS; i++) {
-      total_val = op(static_cast<U>(vals[i]), total_val);
-    }
-
-    in += lsize.x * N_READS;
-  }
-
-  // Separate case for the last set as we close the reduction size   
-  size_t reduction_index = (lid.x + (size_t)lsize.x * r) * N_READS;
-  if(reduction_index < reduction_size) {
-    int max_reads = reduction_size - reduction_index;
-
-    T vals[N_READS]; 
-    for(int i = 0; i < N_READS; i++) {
-      int idx = min(i, max_reads - 1);
-      vals[i] = static_cast<U>(in[idx]);
-    }
-    for(int i = 0; i < N_READS; i++) {
-      T val = i < max_reads ? vals[i] : Op::init;
-      total_val = op(static_cast<U>(val), total_val);
-    }
-  }
+  U total_val = per_thread_row_reduce<T, U, Op, N_READS>(in, reduction_size, out_size, shape, strides, ndim, lsize.x, lid.x, tid.xy);
 
   total_val = op.simd_reduce(total_val);
   
@@ -194,6 +279,53 @@ template <typename T, typename U, typename Op, int N_READS=REDUCE_N_READS>
   }
 }
 
+template <typename T, typename U, typename Op, int N_READS=REDUCE_N_READS>
+[[kernel]] void row_reduce_general_no_atomics(
+    const device T *in [[buffer(0)]],
+    device U *out [[buffer(1)]],
+    const constant size_t& reduction_size [[buffer(2)]],
+    const constant size_t& out_size [[buffer(3)]],
+    const constant int* shape [[buffer(4)]],
+    const constant size_t* strides [[buffer(5)]],
+    const constant int& ndim [[buffer(6)]],
+    uint3 lid [[thread_position_in_threadgroup]],
+    uint3 lsize [[threads_per_threadgroup]],
+    uint3 gsize [[threads_per_grid]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint simd_per_group [[simdgroups_per_threadgroup]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]]) {
+
+  Op op;
+
+  threadgroup U local_vals[simd_size];
+  U total_val = per_thread_row_reduce<T, U, Op, N_READS>(in, reduction_size, out_size, shape, strides, ndim, lsize.x, lid.x, tid.xy);
+
+  // Reduction within simd group - simd_add isn't supported for int64 types
+  for (uint16_t i = simd_size/2; i > 0; i /= 2) {
+    total_val = op(total_val, simd_shuffle_down(total_val, i));
+  }
+
+  // Prepare next level
+  if (simd_lane_id == 0) {
+    local_vals[simd_group_id] = total_val;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  // Reduction within thread group
+  // Only needed if thread group has multiple simd groups
+  if(ceildiv(reduction_size, N_READS) > simd_size) {
+    total_val = lid.x < simd_per_group ? local_vals[lid.x] : op.init;
+    for (uint16_t i = simd_size/2; i > 0; i /= 2) {
+      total_val = op(total_val, simd_shuffle_down(total_val, i));
+    }
+  }
+  // Write row reduce output for threadgroup with 1st thread in thread group
+  if (lid.x == 0) {
+    out[(ceildiv(gsize.y, lsize.y) * tid.x) + tid.y] = total_val;
+  }
+}
+
 #define instantiate_row_reduce_general(name, itype, otype, op) \
   template [[host_name("row_reduce_general_" #name)]] \
   [[kernel]] void row_reduce_general<itype, otype, op>( \
@@ -211,52 +343,59 @@ template <typename T, typename U, typename Op, int N_READS=REDUCE_N_READS>
       uint simd_per_group [[simdgroups_per_threadgroup]],  \
       uint simd_group_id [[simdgroup_index_in_threadgroup]]);
 
+#define instantiate_row_reduce_general_no_atomics(name, itype, otype, op) \
+  template [[host_name("row_reduce_general_no_atomics_" #name)]] \
+  [[kernel]] void row_reduce_general_no_atomics<itype, otype, op>( \
+      const device itype *in [[buffer(0)]],  \
+      device otype *out [[buffer(1)]],  \
+      const constant size_t& reduction_size [[buffer(2)]],  \
+      const constant size_t& out_size [[buffer(3)]],  \
+      const constant int* shape [[buffer(4)]],  \
+      const constant size_t* strides [[buffer(5)]],  \
+      const constant int& ndim [[buffer(6)]],  \
+      uint3 lid [[thread_position_in_threadgroup]],  \
+      uint3 lsize [[threads_per_threadgroup]],  \
+      uint3 gsize [[threads_per_grid]], \
+      uint3 tid [[threadgroup_position_in_grid]],  \
+      uint simd_lane_id [[thread_index_in_simdgroup]],  \
+      uint simd_per_group [[simdgroups_per_threadgroup]],  \
+      uint simd_group_id [[simdgroup_index_in_threadgroup]]);
 
 ///////////////////////////////////////////////////////////////////////////////
 // Column reduce
 ///////////////////////////////////////////////////////////////////////////////
 
 template <typename T, typename U, typename Op, int N_READS = REDUCE_N_READS>
-inline void _contiguous_strided_reduce(
-    const device T *in, 
-    device mlx_atomic<U> *out, 
-    threadgroup U *local_data, 
-    uint in_idx, 
-    uint out_idx, 
-    uint reduction_size, 
-    uint reduction_stride, 
-    uint2 tid, 
-    uint2 lid, 
+inline U _contiguous_strided_reduce(
+    const device T *in,
+    threadgroup U *local_data,
+    uint in_idx,
+    uint reduction_size,
+    uint reduction_stride,
+    uint2 tid,
+    uint2 lid,
     uint2 lsize) {
 
   Op op;
-  T local_vals[N_READS]; 
+  U total_val = Op::init;
 
   uint base_offset = (tid.y * lsize.y + lid.y) * N_READS;
-
-  for(uint r = 0; r < N_READS; r++) {
-    uint offset = base_offset + r; 
-    offset = offset < reduction_size ? offset : reduction_size - 1; 
-    local_vals[r] = in[in_idx + offset * reduction_stride];
-  }
-
-  U total_val = Op::init; 
   for(uint r = 0; r < N_READS && (base_offset + r) < reduction_size; r++) {
-    total_val = op(static_cast<U>(total_val), local_vals[r]); 
+    uint offset = base_offset + r;
+    total_val = op(static_cast<U>(total_val), in[in_idx + offset * reduction_stride]);
   }
-  local_data[lsize.y * lid.x + lid.y] = total_val; 
-
+  local_data[lsize.y * lid.x + lid.y] = total_val;
   threadgroup_barrier(mem_flags::mem_threadgroup);
 
+  U val = Op::init;
   if(lid.y == 0) {
-    U val = op.init; 
-
+    // Perform reduction across columns in thread group
     for(uint i = 0; i < lsize.y; i++) {
-      val = op(val, local_data[lsize.y * lid.x + i]); 
+      val = op(val, local_data[lsize.y * lid.x + i]);
     }
-
-    op.atomic_update(out, val, out_idx);
   }
+
+  return val;
 }
 
 template <typename T, typename U, typename Op, int N_READS = REDUCE_N_READS>
@@ -265,13 +404,13 @@ template <typename T, typename U, typename Op, int N_READS = REDUCE_N_READS>
     device mlx_atomic<U> *out [[buffer(1)]],
     const constant size_t& reduction_size [[buffer(2)]],
     const constant size_t& reduction_stride [[buffer(3)]],
-    const constant size_t& out_size [[buffer(4)]], 
+    const constant size_t& out_size [[buffer(4)]],
     const constant int* shape [[buffer(5)]],
     const constant size_t* strides [[buffer(6)]],
     const constant int& ndim [[buffer(7)]],
     threadgroup U *local_data [[threadgroup(0)]],
-    uint3 tid [[threadgroup_position_in_grid]], 
-    uint3 lid [[thread_position_in_threadgroup]], 
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint3 lid [[thread_position_in_threadgroup]],
     uint3 lsize [[threads_per_threadgroup]]) {
   auto out_idx = tid.x * lsize.x + lid.x;
   auto in_idx = elem_to_loc(
@@ -281,18 +420,66 @@ template <typename T, typename U, typename Op, int N_READS = REDUCE_N_READS>
     ndim
   );
 
+  Op op;
   if(out_idx < out_size) {
-    _contiguous_strided_reduce<T, U, Op, N_READS>(
-      in, 
-      out, 
-      local_data, 
-      in_idx,
-      out_idx, 
-      reduction_size, 
-      reduction_stride, 
-      tid.xy, 
-      lid.xy, 
-      lsize.xy); 
+    U val = _contiguous_strided_reduce<T, U, Op, N_READS>(
+              in,
+              local_data,
+              in_idx,
+              reduction_size,
+              reduction_stride,
+              tid.xy,
+              lid.xy,
+              lsize.xy);
+
+    // Write out reduction results generated by threadgroups working on specific output element, contiguously.
+    if (lid.y == 0) {
+      op.atomic_update(out, val, out_idx);
+    }
+  }
+}
+
+template <typename T, typename U, typename Op, int N_READS = REDUCE_N_READS>
+[[kernel]] void col_reduce_general_no_atomics(
+    const device T *in [[buffer(0)]],
+    device U *out [[buffer(1)]],
+    const constant size_t& reduction_size [[buffer(2)]],
+    const constant size_t& reduction_stride [[buffer(3)]],
+    const constant size_t& out_size [[buffer(4)]],
+    const constant int* shape [[buffer(5)]],
+    const constant size_t* strides [[buffer(6)]],
+    const constant int& ndim [[buffer(7)]],
+    threadgroup U *local_data [[threadgroup(0)]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint3 lid [[thread_position_in_threadgroup]],
+    uint3 gid [[thread_position_in_grid]],
+    uint3 lsize [[threads_per_threadgroup]],
+    uint3 gsize [[threads_per_grid]]) {
+  auto out_idx = tid.x * lsize.x + lid.x;
+  auto in_idx = elem_to_loc(
+    out_idx + tid.z * out_size,
+    shape,
+    strides,
+    ndim
+  );
+
+  if(out_idx < out_size) {
+    U val = _contiguous_strided_reduce<T, U, Op, N_READS>(
+              in,
+              local_data,
+              in_idx,
+              reduction_size,
+              reduction_stride,
+              tid.xy,
+              lid.xy,
+              lsize.xy);
+
+    // Write out reduction results generated by threadgroups working on specific output element, contiguously.
+    if (lid.y == 0) {
+      uint tgsize_y = ceildiv(gsize.y, lsize.y);
+      uint tgsize_z = ceildiv(gsize.z, lsize.z);
+      out[tgsize_y * tgsize_z * gid.x + tgsize_y * tid.z + tid.y] = val;
+    }
   }
 }
 
@@ -312,6 +499,23 @@ template <typename T, typename U, typename Op, int N_READS = REDUCE_N_READS>
       uint3 lid [[thread_position_in_threadgroup]], \
       uint3 lsize [[threads_per_threadgroup]]);
 
+#define instantiate_col_reduce_general_no_atomics(name, itype, otype, op) \
+  template [[host_name("col_reduce_general_no_atomics_" #name)]] \
+  [[kernel]] void col_reduce_general_no_atomics<itype, otype, op>( \
+      const device itype *in [[buffer(0)]], \
+      device otype *out [[buffer(1)]], \
+      const constant size_t& reduction_size [[buffer(2)]], \
+      const constant size_t& reduction_stride [[buffer(3)]], \
+      const constant size_t& out_size [[buffer(4)]], \
+      const constant int* shape [[buffer(5)]],  \
+      const constant size_t* strides [[buffer(6)]],  \
+      const constant int& ndim [[buffer(7)]],  \
+      threadgroup otype *local_data [[threadgroup(0)]], \
+      uint3 tid [[threadgroup_position_in_grid]], \
+      uint3 lid [[thread_position_in_threadgroup]], \
+      uint3 gid [[thread_position_in_grid]], \
+      uint3 lsize [[threads_per_threadgroup]], \
+      uint3 gsize [[threads_per_grid]]);
 
 ///////////////////////////////////////////////////////////////////////////////
 // Instantiations
@@ -322,6 +526,15 @@ template <typename T, typename U, typename Op, int N_READS = REDUCE_N_READS>
   instantiate_row_reduce_general(name, itype, otype, op) \
   instantiate_col_reduce_general(name, itype, otype, op)
 
+#define instantiate_reduce_no_atomics(name, itype, otype, op) \
+  instantiate_all_reduce_no_atomics(name, itype, otype, op) \
+  instantiate_row_reduce_general_no_atomics(name, itype, otype, op) \
+  instantiate_col_reduce_general_no_atomics(name, itype, otype, op)
+
+#define instantiate_same_reduce_no_atomics(name, tname, type, op) \
+  instantiate_init_reduce(name ##tname, type, op<type>) \
+  instantiate_reduce_no_atomics(name ##tname, type, type, op<type>)
+
 #define instantiate_same_reduce(name, tname, type, op) \
   instantiate_init_reduce(name ##tname, type, op<type>) \
   instantiate_reduce(name ##tname, type, type, op<type>)
@@ -353,6 +566,9 @@ instantiate_same_reduce(sum, int32, int32_t, Sum)
 instantiate_same_reduce(sum, float16, half, Sum)
 instantiate_same_reduce(sum, float32, float, Sum)
 
+instantiate_same_reduce_no_atomics(sum, int64, int64_t, Sum)
+instantiate_same_reduce_no_atomics(sum, uint64, uint64_t, Sum)
+
 instantiate_same_reduce(prod, uint8, uint8_t, Prod)
 instantiate_same_reduce(prod, uint16, uint16_t, Prod)
 instantiate_same_reduce(prod, uint32, uint32_t, Prod)
@@ -362,6 +578,9 @@ instantiate_same_reduce(prod, int32, int32_t, Prod)
 instantiate_same_reduce(prod, float16, half, Prod)
 instantiate_same_reduce(prod, float32, float, Prod)
 
+instantiate_same_reduce_no_atomics(prod, int64, int64_t, Prod)
+instantiate_same_reduce_no_atomics(prod, uint64, uint64_t, Prod)
+
 instantiate_same_reduce(sum, bfloat16, bfloat16_t, Sum)
 instantiate_same_reduce(prod, bfloat16, bfloat16_t, Prod)
 
@@ -381,6 +600,9 @@ instantiate_same_reduce(min_, int32, int32_t, Min)
 instantiate_same_reduce(min_, float16, half, Min)
 instantiate_same_reduce(min_, float32, float, Min)
 
+instantiate_same_reduce_no_atomics(min_, int64, int64_t, Min)
+instantiate_same_reduce_no_atomics(min_, uint64, uint64_t, Min)
+
 instantiate_same_reduce(max_, uint8, uint8_t, Max)
 instantiate_same_reduce(max_, uint16, uint16_t, Max)
 instantiate_same_reduce(max_, uint32, uint32_t, Max)
@@ -390,5 +612,8 @@ instantiate_same_reduce(max_, int32, int32_t, Max)
 instantiate_same_reduce(max_, float16, half, Max)
 instantiate_same_reduce(max_, float32, float, Max)
 
+instantiate_same_reduce_no_atomics(max_, int64, int64_t, Max)
+instantiate_same_reduce_no_atomics(max_, uint64, uint64_t, Max)
+
 instantiate_same_reduce(min_, bfloat16, bfloat16_t, Min)
 instantiate_same_reduce(max_, bfloat16, bfloat16_t, Max)

mlx/backend/metal/kernels/utils.h

@@ -235,12 +235,42 @@ inline size_t ceildiv(size_t N, size_t M) {
 // https://docs.oracle.com/cd/E19957-01/806-3568/ncg_goldberg.html#1202
 inline float log1p(float x) {
   float xp1 = 1.0f + x;
-  return (xp1 == 1.0f) ? x : x * (metal::log(xp1) / (xp1 - 1.0f));
+  if (xp1 == Limits<float>::max) {
+    return Limits<float>::max;
+  }
+  if (xp1 == 1.0f) {
+    return x;
+  }
+
+  return x * (metal::log(xp1) / (xp1 - 1.0f));
 }
 
 inline bfloat16_t log1p(bfloat16_t x) {
   float xp1 = 1.0f + static_cast<float>(x);
-  bfloat16_t ret =
-      (xp1 == 1.0f) ? x : bfloat16_t(x * (metal::log(xp1) / (xp1 - 1.0f)));
-  return ret;
+  if (xp1 == Limits<float>::max) {
+    return Limits<bfloat16_t>::max;
+  }
+  if (xp1 == 1.0f) {
+    return x;
+  }
+
+  return bfloat16_t(x * (metal::log(xp1) / (xp1 - 1.0f)));
 }
+
+///////////////////////////////////////////////////////////////////////////////
+// SIMD shuffle ops
+///////////////////////////////////////////////////////////////////////////////
+
+inline uint64_t simd_shuffle_down(uint64_t data, uint16_t delta) {
+  return as_type<uint64_t>(
+      metal::simd_shuffle_down(as_type<uint2>(data), delta));
+}
+
+inline int64_t simd_shuffle_down(int64_t data, uint16_t delta) {
+  return as_type<int64_t>(
+      metal::simd_shuffle_down(as_type<uint2>(data), delta));
+}
+
+inline bool simd_shuffle_down(bool data, uint16_t delta) {
+  return simd_shuffle_down(static_cast<uint32_t>(data), delta);
+}

mlx/backend/metal/matmul.cpp

@@ -1,4 +1,4 @@
-// Copyright © 2023 Apple Inc.
+// Copyright © 2023-2024 Apple Inc.
 
 #include <algorithm>
 #include <cassert>
@@ -615,7 +615,7 @@ void Matmul::eval_gpu(const std::vector<array>& inputs, array& out) {
 }
 
 void AddMM::eval_gpu(const std::vector<array>& inputs, array& out) {
-  assert(inputs.size() == 2);
+  assert(inputs.size() == 3);
   if (!is_floating_point(out.dtype())) {
     throw std::runtime_error(
         "[matmul] Does not yet support non-floating point types.");

mlx/backend/metal/metal.cpp

@@ -64,14 +64,23 @@ std::function<void()> make_task(
     auto command_buffer = increment_command_buffer(s);
     auto outputs = arr.outputs();
     arr.primitive().eval_gpu(arr.inputs(), outputs);
+    std::vector<std::shared_ptr<array::Data>> buffers;
+    for (auto& in : arr.inputs()) {
+      buffers.push_back(in.data_shared_ptr());
+    }
+    for (auto& s : arr.siblings()) {
+      buffers.push_back(s.data_shared_ptr());
+    }
+    if (!arr.is_tracer()) {
+      arr.detach();
+    }
+
     if (p) {
       metal::device(s.device).end_encoding(s.index);
       scheduler::notify_new_task(s);
       command_buffer->addCompletedHandler(
-          [s, arr, p = std::move(p)](MTL::CommandBuffer* cbuf) mutable {
-            if (!arr.is_tracer()) {
-              arr.detach();
-            }
+          [s, buffers = std::move(buffers), p = std::move(p)](
+              MTL::CommandBuffer* cbuf) {
             p->set_value();
             scheduler::notify_task_completion(s);
             check_error(cbuf);
@@ -79,10 +88,7 @@ std::function<void()> make_task(
       metal::device(s.device).commit_command_buffer(s.index);
     } else {
       command_buffer->addCompletedHandler(
-          [s, arr](MTL::CommandBuffer* cbuf) mutable {
-            if (!arr.is_tracer()) {
-              arr.detach();
-            }
+          [s, buffers = std::move(buffers)](MTL::CommandBuffer* cbuf) {
             check_error(cbuf);
           });
     }

mlx/backend/metal/primitives.cpp

@@ -1,4 +1,4 @@
-// Copyright © 2023 Apple Inc.
+// Copyright © 2023-2024 Apple Inc.
 
 #include <algorithm>
 #include <cassert>
@@ -27,8 +27,8 @@ void binary_op(
   auto& a = inputs[0];
   auto& b = inputs[1];
   auto bopt = get_binary_op_type(a, b);
-  set_binary_op_output_data(a, b, outputs[0], bopt);
-  set_binary_op_output_data(a, b, outputs[1], bopt);
+  set_binary_op_output_data(a, b, outputs[0], bopt, true);
+  set_binary_op_output_data(a, b, outputs[1], bopt, true);
 
   auto& out = outputs[0];
   if (out.size() == 0) {
@@ -60,7 +60,7 @@ void binary_op(
       break;
   }
   kname << op << type_to_name(a);
-  if (bopt == General && out.ndim() <= MAX_BINARY_SPECIALIZED_DIMS) {
+  if (bopt == General && shape.size() <= MAX_BINARY_SPECIALIZED_DIMS) {
     kname << "_" << shape.size();
   }
 
@@ -69,8 +69,14 @@ void binary_op(
   auto kernel = d.get_kernel(kname.str());
   auto compute_encoder = d.get_command_encoder(s.index);
   compute_encoder->setComputePipelineState(kernel);
-  set_array_buffer(compute_encoder, a, 0);
-  set_array_buffer(compute_encoder, b, 1);
+  // - If a is donated it goes to the first output
+  // - If b is donated it goes to the first output if a was not donated
+  //   otherwise it goes to the second output
+  bool donate_a = a.data_shared_ptr() == nullptr;
+  bool donate_b = b.data_shared_ptr() == nullptr;
+  set_array_buffer(compute_encoder, donate_a ? outputs[0] : a, 0);
+  set_array_buffer(
+      compute_encoder, donate_b ? (donate_a ? outputs[1] : outputs[0]) : b, 1);
   set_array_buffer(compute_encoder, outputs[0], 2);
   set_array_buffer(compute_encoder, outputs[1], 3);
 
@@ -122,7 +128,7 @@ void binary_op(
   auto& a = inputs[0];
   auto& b = inputs[1];
   auto bopt = get_binary_op_type(a, b);
-  set_binary_op_output_data(a, b, out, bopt);
+  set_binary_op_output_data(a, b, out, bopt, true);
   if (out.size() == 0) {
     return;
   }
@@ -152,7 +158,7 @@ void binary_op(
       break;
   }
   kname << op << type_to_name(a);
-  if (bopt == General && out.ndim() <= MAX_BINARY_SPECIALIZED_DIMS) {
+  if (bopt == General && shape.size() <= MAX_BINARY_SPECIALIZED_DIMS) {
     kname << "_" << shape.size();
   }
 
@@ -161,8 +167,10 @@ void binary_op(
   auto kernel = d.get_kernel(kname.str());
   auto compute_encoder = d.get_command_encoder(s.index);
   compute_encoder->setComputePipelineState(kernel);
-  set_array_buffer(compute_encoder, a, 0);
-  set_array_buffer(compute_encoder, b, 1);
+  bool donate_a = a.data_shared_ptr() == nullptr;
+  bool donate_b = b.data_shared_ptr() == nullptr;
+  set_array_buffer(compute_encoder, donate_a ? out : a, 0);
+  set_array_buffer(compute_encoder, donate_b ? out : b, 1);
   set_array_buffer(compute_encoder, out, 2);
 
   if (bopt == General) {
@@ -212,11 +220,15 @@ void unary_op(
   auto& in = inputs[0];
   bool contig = in.flags().contiguous;
   if (contig) {
-    out.set_data(
-        allocator::malloc_or_wait(in.data_size() * out.itemsize()),
-        in.data_size(),
-        in.strides(),
-        in.flags());
+    if (in.is_donatable() && in.itemsize() == out.itemsize()) {
+      out.move_shared_buffer(in);
+    } else {
+      out.set_data(
+          allocator::malloc_or_wait(in.data_size() * out.itemsize()),
+          in.data_size(),
+          in.strides(),
+          in.flags());
+    }
   } else {
     out.set_data(allocator::malloc_or_wait(out.nbytes()));
   }
@@ -240,7 +252,8 @@ void unary_op(
 
   auto compute_encoder = d.get_command_encoder(s.index);
   compute_encoder->setComputePipelineState(kernel);
-  set_array_buffer(compute_encoder, in, 0);
+  set_array_buffer(
+      compute_encoder, in.data_shared_ptr() == nullptr ? out : in, 0);
   set_array_buffer(compute_encoder, out, 1);
   if (!contig) {
     compute_encoder->setBytes(in.shape().data(), in.ndim() * sizeof(int), 2);
@@ -473,6 +486,18 @@ void Cosh::eval_gpu(const std::vector<array>& inputs, array& out) {
   unary_op(inputs, out, "cosh");
 }
 
+void CustomVJP::eval_gpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  eval(inputs, outputs);
+}
+
+void Depends::eval_gpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  eval(inputs, outputs);
+}
+
 void Divide::eval_gpu(const std::vector<array>& inputs, array& out) {
   binary_op(inputs, out, "div");
 }
@@ -769,4 +794,10 @@ void Transpose::eval_gpu(const std::vector<array>& inputs, array& out) {
   eval(inputs, out);
 }
 
+void QRF::eval_gpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  throw std::runtime_error("[QRF::eval_gpu] Metal QR factorization NYI.");
+}
+
 } // namespace mlx::core

mlx/backend/metal/quantized.cpp

@@ -1,4 +1,4 @@
-// Copyright © 2023 Apple Inc.
+// Copyright © 2023-2024 Apple Inc.
 
 #include <cassert>
 

mlx/backend/metal/reduce.cpp

@@ -28,35 +28,40 @@ inline auto safe_divup(size_t n, size_t m) {
   return safe_div(n, m) * m;
 }
 
+inline bool is_64b_int(Dtype dtype) {
+  return dtype == int64 || dtype == uint64;
+}
+
 // All Reduce
 void all_reduce_dispatch(
     const array& in,
     array& out,
     const std::string& op_name,
     MTL::ComputeCommandEncoder* compute_encoder,
-    metal::Device& d) {
-  // Get kernel and encode buffers
-  size_t in_size = in.size();
-  auto kernel = d.get_kernel("all_reduce_" + op_name + type_to_name(in));
+    metal::Device& d,
+    const Stream& s) {
+  Dtype out_dtype = out.dtype();
+  bool is_out_64b_int = is_64b_int(out_dtype);
+  auto kernel = (is_out_64b_int)
+      ? d.get_kernel("all_reduce_no_atomics_" + op_name + type_to_name(in))
+      : d.get_kernel("all_reduce_" + op_name + type_to_name(in));
 
   compute_encoder->setComputePipelineState(kernel);
-  set_array_buffer(compute_encoder, in, 0);
-  set_array_buffer(compute_encoder, out, 1);
-  compute_encoder->setBytes(&in_size, sizeof(size_t), 2);
-
-  // Set grid dimensions
 
   // We make sure each thread has enough to do by making it read in
   // at least n_reads inputs
   int n_reads = REDUCE_N_READS;
+  size_t in_size = in.size();
 
   // mod_in_size gives us the groups of n_reads needed to go over the entire
   // input
   uint mod_in_size = (in_size + n_reads - 1) / n_reads;
-
   NS::UInteger thread_group_size = kernel->maxTotalThreadsPerThreadgroup();
   thread_group_size =
       mod_in_size > thread_group_size ? thread_group_size : mod_in_size;
+  uint simd_size = kernel->threadExecutionWidth();
+  thread_group_size =
+      ((thread_group_size + simd_size - 1) / simd_size) * simd_size;
 
   // If the number of thread groups needed exceeds 1024, we reuse threads groups
   uint n_thread_groups = safe_div(mod_in_size, thread_group_size);
@@ -66,7 +71,52 @@ void all_reduce_dispatch(
   MTL::Size group_dims = MTL::Size(thread_group_size, 1, 1);
   MTL::Size grid_dims = MTL::Size(nthreads, 1, 1);
 
-  compute_encoder->dispatchThreads(grid_dims, group_dims);
+  // Encode buffers and dispatch
+  if (is_out_64b_int == false || n_thread_groups == 1) {
+    set_array_buffer(compute_encoder, in, 0);
+    set_array_buffer(compute_encoder, out, 1);
+    compute_encoder->setBytes(&in_size, sizeof(size_t), 2);
+    compute_encoder->dispatchThreads(grid_dims, group_dims);
+
+  } else {
+    // Allocate intermediate array to store partial reduction results
+    size_t intermediate_size = n_thread_groups;
+    array intermediate =
+        array({static_cast<int>(intermediate_size)}, out_dtype, nullptr, {});
+    intermediate.set_data(allocator::malloc_or_wait(intermediate.nbytes()));
+    std::vector<array> intermediates = {intermediate};
+
+    // First dispatch
+    set_array_buffer(compute_encoder, in, 0);
+    set_array_buffer(compute_encoder, intermediate, 1);
+    compute_encoder->setBytes(&in_size, sizeof(size_t), 2);
+    compute_encoder->dispatchThreads(grid_dims, group_dims);
+
+    // Second pass to reduce intermediate reduction results written to DRAM
+    set_array_buffer(compute_encoder, intermediate, 0);
+    set_array_buffer(compute_encoder, out, 1);
+    compute_encoder->setBytes(&intermediate_size, sizeof(size_t), 2);
+
+    mod_in_size = (intermediate_size + n_reads - 1) / n_reads;
+
+    thread_group_size = kernel->maxTotalThreadsPerThreadgroup();
+    thread_group_size =
+        mod_in_size > thread_group_size ? thread_group_size : mod_in_size;
+    thread_group_size =
+        ((thread_group_size + simd_size - 1) / simd_size) * simd_size;
+
+    // If the number of thread groups needed exceeds 1024, we reuse threads
+    // groups
+    nthreads = thread_group_size;
+    group_dims = MTL::Size(thread_group_size, 1, 1);
+    grid_dims = MTL::Size(nthreads, 1, 1);
+    compute_encoder->dispatchThreads(grid_dims, group_dims);
+
+    d.get_command_buffer(s.index)->addCompletedHandler(
+        [intermediates](MTL::CommandBuffer*) mutable {
+          intermediates.clear();
+        });
+  }
 }
 
 void row_reduce_general_dispatch(
@@ -76,22 +126,31 @@ void row_reduce_general_dispatch(
     const ReductionPlan& plan,
     const std::vector<int>& axes,
     MTL::ComputeCommandEncoder* compute_encoder,
-    metal::Device& d) {
-  auto kernel =
-      d.get_kernel("row_reduce_general_" + op_name + type_to_name(in));
+    metal::Device& d,
+    const Stream& s) {
+  Dtype out_dtype = out.dtype();
+  bool is_out_64b_int = is_64b_int(out_dtype);
+  auto kernel = (is_out_64b_int)
+      ? d.get_kernel(
+            "row_reduce_general_no_atomics_" + op_name + type_to_name(in))
+      : d.get_kernel("row_reduce_general_" + op_name + type_to_name(in));
+
+  compute_encoder->setComputePipelineState(kernel);
 
   // Prepare the arguments for the kernel
   int n_reads = REDUCE_N_READS;
   size_t reduction_size = plan.shape.back();
-  size_t out_size = out.size();
   auto shape = plan.shape;
   auto strides = plan.strides;
+
   shape.pop_back();
   strides.pop_back();
+
   size_t non_row_reductions = 1;
   for (auto s : shape) {
     non_row_reductions *= static_cast<size_t>(s);
   }
+  size_t out_size = out.size();
   auto [rem_shape, rem_strides] = shapes_without_reduction_axes(in, axes);
   for (auto s : rem_shape) {
     shape.push_back(s);
@@ -101,16 +160,6 @@ void row_reduce_general_dispatch(
   }
   int ndim = shape.size();
 
-  // Set the arguments for the kernel
-  compute_encoder->setComputePipelineState(kernel);
-  set_array_buffer(compute_encoder, in, 0);
-  set_array_buffer(compute_encoder, out, 1);
-  compute_encoder->setBytes(&reduction_size, sizeof(size_t), 2);
-  compute_encoder->setBytes(&out_size, sizeof(size_t), 3);
-  compute_encoder->setBytes(shape.data(), shape.size() * sizeof(int), 4);
-  compute_encoder->setBytes(strides.data(), strides.size() * sizeof(size_t), 5);
-  compute_encoder->setBytes(&ndim, sizeof(int), 6);
-
   // Each thread group is responsible for 1 output
   NS::UInteger thread_group_size = kernel->maxTotalThreadsPerThreadgroup();
   thread_group_size =
@@ -127,7 +176,88 @@ void row_reduce_general_dispatch(
   MTL::Size grid_dims = MTL::Size(n_threads, non_row_reductions, 1);
   MTL::Size group_dims = MTL::Size(thread_group_size, 1, 1);
 
-  compute_encoder->dispatchThreads(grid_dims, group_dims);
+  if (is_out_64b_int == false || non_row_reductions == 1) {
+    // Set the arguments for the kernel
+    set_array_buffer(compute_encoder, in, 0);
+    set_array_buffer(compute_encoder, out, 1);
+    compute_encoder->setBytes(&reduction_size, sizeof(size_t), 2);
+    compute_encoder->setBytes(&out_size, sizeof(size_t), 3);
+    compute_encoder->setBytes(shape.data(), shape.size() * sizeof(int), 4);
+    compute_encoder->setBytes(
+        strides.data(), strides.size() * sizeof(size_t), 5);
+    compute_encoder->setBytes(&ndim, sizeof(int), 6);
+    compute_encoder->dispatchThreads(grid_dims, group_dims);
+
+  } else {
+    // Allocate intermediate array to store partial reduction results
+    array intermediate = array(
+        {static_cast<int>(out.size()), static_cast<int>(non_row_reductions)},
+        out_dtype,
+        nullptr,
+        {});
+    intermediate.set_data(allocator::malloc_or_wait(intermediate.nbytes()));
+    std::vector<array> intermediates = {intermediate};
+
+    // Set the arguments for the kernel
+    set_array_buffer(compute_encoder, in, 0);
+    set_array_buffer(compute_encoder, intermediate, 1);
+    compute_encoder->setBytes(&reduction_size, sizeof(size_t), 2);
+    compute_encoder->setBytes(&out_size, sizeof(size_t), 3);
+    compute_encoder->setBytes(shape.data(), shape.size() * sizeof(int), 4);
+    compute_encoder->setBytes(
+        strides.data(), strides.size() * sizeof(size_t), 5);
+    compute_encoder->setBytes(&ndim, sizeof(int), 6);
+    compute_encoder->dispatchThreads(grid_dims, group_dims);
+
+    // Set up second dispatch
+    reduction_size = non_row_reductions;
+    out_size = 1;
+
+    // Shape of axes that aren't participating in reduction remains unchanged.
+    std::vector<int> new_shape = rem_shape;
+
+    // Update their strides since they'll be different post partial reduction in
+    // first compute dispatch.
+    std::vector<size_t> new_strides = rem_strides;
+    new_strides.back() = reduction_size;
+    for (int i = new_shape.size() - 2; i >= 0; i--) {
+      new_strides[i] = new_shape[i + 1] * new_strides[i + 1];
+    }
+    ndim = new_shape.size();
+
+    // Set the arguments for the kernel
+    set_array_buffer(compute_encoder, intermediate, 0);
+    set_array_buffer(compute_encoder, out, 1);
+    compute_encoder->setBytes(&reduction_size, sizeof(size_t), 2);
+    compute_encoder->setBytes(&out_size, sizeof(size_t), 3);
+    compute_encoder->setBytes(
+        new_shape.data(), new_shape.size() * sizeof(int), 4);
+    compute_encoder->setBytes(
+        new_strides.data(), new_strides.size() * sizeof(size_t), 5);
+    compute_encoder->setBytes(&ndim, sizeof(int), 6);
+
+    // Each thread group is responsible for 1 output
+    thread_group_size = kernel->maxTotalThreadsPerThreadgroup();
+    thread_group_size =
+        std::min((reduction_size + n_reads - 1) / n_reads, thread_group_size);
+
+    // Align thread group size with simd_size
+    thread_group_size =
+        (thread_group_size + simd_size - 1) / simd_size * simd_size;
+    assert(thread_group_size <= kernel->maxTotalThreadsPerThreadgroup());
+
+    // Launch enough thread groups for each output
+    n_threads = thread_group_size;
+    grid_dims = MTL::Size(n_threads, out.size(), 1);
+    group_dims = MTL::Size(thread_group_size, 1, 1);
+
+    compute_encoder->dispatchThreads(grid_dims, group_dims);
+
+    d.get_command_buffer(s.index)->addCompletedHandler(
+        [intermediates](MTL::CommandBuffer*) mutable {
+          intermediates.clear();
+        });
+  }
 }
 
 void strided_reduce_general_dispatch(
@@ -137,9 +267,16 @@ void strided_reduce_general_dispatch(
     const ReductionPlan& plan,
     const std::vector<int>& axes,
     MTL::ComputeCommandEncoder* compute_encoder,
-    metal::Device& d) {
-  auto kernel =
-      d.get_kernel("col_reduce_general_" + op_name + type_to_name(in));
+    metal::Device& d,
+    const Stream& s) {
+  Dtype out_dtype = out.dtype();
+  bool is_out_64b_int = is_64b_int(out_dtype);
+  auto kernel = (is_out_64b_int)
+      ? d.get_kernel(
+            "col_reduce_general_no_atomics_" + op_name + type_to_name(in))
+      : d.get_kernel("col_reduce_general_" + op_name + type_to_name(in));
+
+  compute_encoder->setComputePipelineState(kernel);
 
   // Prepare the arguments for the kernel
   size_t reduction_size = plan.shape.back();
@@ -162,19 +299,7 @@ void strided_reduce_general_dispatch(
   }
   int ndim = shape.size();
 
-  // Set the arguments for the kernel
-  compute_encoder->setComputePipelineState(kernel);
-  set_array_buffer(compute_encoder, in, 0);
-  set_array_buffer(compute_encoder, out, 1);
-  compute_encoder->setBytes(&reduction_size, sizeof(size_t), 2);
-  compute_encoder->setBytes(&reduction_stride, sizeof(size_t), 3);
-  compute_encoder->setBytes(&out_size, sizeof(size_t), 4);
-  compute_encoder->setBytes(shape.data(), shape.size() * sizeof(int), 5);
-  compute_encoder->setBytes(strides.data(), strides.size() * sizeof(size_t), 6);
-  compute_encoder->setBytes(&ndim, sizeof(int), 7);
-
   // Select block dimensions
-
   // Each thread reads 16 inputs to give it more work
   uint n_inputs_per_thread = REDUCE_N_READS;
   uint n_threads_per_output =
@@ -183,14 +308,22 @@ void strided_reduce_general_dispatch(
   // We spread outputs over the x dimension and inputs over the y dimension
   // Threads with the same lid.x in a given threadgroup work on the same
   // output and each thread in the y dimension accumulates for that output
+
+  // Threads with same lid.x, i.e. each column of threads work on same output
   uint threadgroup_dim_x = std::min(out_size, 128ul);
+
+  // Number of threads along y, is dependent on number of reductions needed.
   uint threadgroup_dim_y =
       kernel->maxTotalThreadsPerThreadgroup() / threadgroup_dim_x;
   threadgroup_dim_y = std::min(n_threads_per_output, threadgroup_dim_y);
 
+  // Derive number of thread groups along x, based on how many threads we need
+  // along x
   uint n_threadgroups_x =
       (out_size + threadgroup_dim_x - 1) / threadgroup_dim_x;
 
+  // Derive number of thread groups along y based on how many threads we need
+  // along y
   uint n_threadgroups_y =
       (n_threads_per_output + threadgroup_dim_y - 1) / threadgroup_dim_y;
 
@@ -199,18 +332,122 @@ void strided_reduce_general_dispatch(
       MTL::Size(n_threadgroups_x, n_threadgroups_y, non_col_reductions);
   MTL::Size group_dims = MTL::Size(threadgroup_dim_x, threadgroup_dim_y, 1);
 
-  // We set shared memory to be exploited here for reductions within a
-  // threadgroup - each thread must be able to update its accumulated output
-  // Note: Each threadgroup should have 32kB of data in threadgroup memory
-  //       and threadgroup_dim_x * threadgroup_dim_y <= 1024 by design
-  //       This should be fine for floats, but we might need to revisit
-  //       if we ever come to doubles. In that case, we should also cut
-  //       down the number of threads we launch in a threadgroup
-  compute_encoder->setThreadgroupMemoryLength(
-      safe_divup(threadgroup_dim_x * threadgroup_dim_y * out.itemsize(), 16),
-      0);
+  if (is_out_64b_int == false) {
+    // Set the arguments for the kernel
+    set_array_buffer(compute_encoder, in, 0);
+    set_array_buffer(compute_encoder, out, 1);
+    compute_encoder->setBytes(&reduction_size, sizeof(size_t), 2);
+    compute_encoder->setBytes(&reduction_stride, sizeof(size_t), 3);
+    compute_encoder->setBytes(&out_size, sizeof(size_t), 4);
+    compute_encoder->setBytes(shape.data(), shape.size() * sizeof(int), 5);
+    compute_encoder->setBytes(
+        strides.data(), strides.size() * sizeof(size_t), 6);
+    compute_encoder->setBytes(&ndim, sizeof(int), 7);
 
-  compute_encoder->dispatchThreadgroups(grid_dims, group_dims);
+    // We set shared memory to be exploited here for reductions within a
+    // threadgroup - each thread must be able to update its accumulated output
+    // Note: Each threadgroup should have 32kB of data in threadgroup memory
+    //       and threadgroup_dim_x * threadgroup_dim_y <= 1024 by design
+    //       This should be fine for floats, but we might need to revisit
+    //       if we ever come to doubles. In that case, we should also cut
+    //       down the number of threads we launch in a threadgroup
+    compute_encoder->setThreadgroupMemoryLength(
+        safe_divup(threadgroup_dim_x * threadgroup_dim_y * out.itemsize(), 16),
+        0);
+    compute_encoder->dispatchThreadgroups(grid_dims, group_dims);
+
+  } else {
+    // Allocate intermediate array to store reduction results from all thread
+    // groups
+    array intermediate = array(
+        {static_cast<int>(out.size()),
+         static_cast<int>(n_threadgroups_y * non_col_reductions)},
+        out_dtype,
+        nullptr,
+        {});
+    intermediate.set_data(allocator::malloc_or_wait(intermediate.nbytes()));
+    std::vector<array> intermediates = {intermediate};
+
+    // Set the arguments for the kernel
+    set_array_buffer(compute_encoder, in, 0);
+    set_array_buffer(compute_encoder, intermediate, 1);
+    compute_encoder->setBytes(&reduction_size, sizeof(size_t), 2);
+    compute_encoder->setBytes(&reduction_stride, sizeof(size_t), 3);
+    compute_encoder->setBytes(&out_size, sizeof(size_t), 4);
+    compute_encoder->setBytes(shape.data(), shape.size() * sizeof(int), 5);
+    compute_encoder->setBytes(
+        strides.data(), strides.size() * sizeof(size_t), 6);
+    compute_encoder->setBytes(&ndim, sizeof(int), 7);
+
+    // We set shared memory to be exploited here for reductions within a
+    // threadgroup - each thread must be able to update its accumulated output
+    // Note: Each threadgroup should have 32kB of data in threadgroup memory
+    //       and threadgroup_dim_x * threadgroup_dim_y <= 1024 by design
+    //       This should be fine for floats, but we might need to revisit
+    //       if we ever come to doubles. In that case, we should also cut
+    //       down the number of threads we launch in a threadgroup
+    compute_encoder->setThreadgroupMemoryLength(
+        safe_divup(threadgroup_dim_x * threadgroup_dim_y * out.itemsize(), 16),
+        0);
+    compute_encoder->dispatchThreadgroups(grid_dims, group_dims);
+
+    // Perform second pass of reductions
+    // Reduce results of threadgroups along y, z from first pass, that
+    // collectively work on each output element.
+    reduction_size = n_threadgroups_y * non_col_reductions;
+    out_size = 1;
+
+    // Shape of axes that aren't participating in reduction remains unchanged.
+    std::vector<int> new_shape = rem_shape;
+
+    // Update their strides since they'll be different after a partial reduction
+    // post first compute dispatch.
+    std::vector<size_t> new_strides = rem_strides;
+    new_strides.back() = reduction_size;
+    for (int i = new_shape.size() - 2; i >= 0; i--) {
+      new_strides[i] = new_shape[i + 1] * new_strides[i + 1];
+    }
+    ndim = new_shape.size();
+
+    auto row_reduce_kernel = d.get_kernel(
+        "row_reduce_general_no_atomics_" + op_name +
+        type_to_name(intermediate));
+    compute_encoder->setComputePipelineState(row_reduce_kernel);
+    set_array_buffer(compute_encoder, intermediate, 0);
+    set_array_buffer(compute_encoder, out, 1);
+    compute_encoder->setBytes(&reduction_size, sizeof(size_t), 2);
+    compute_encoder->setBytes(&out_size, sizeof(size_t), 3);
+    compute_encoder->setBytes(
+        new_shape.data(), new_shape.size() * sizeof(int), 4);
+    compute_encoder->setBytes(
+        new_strides.data(), new_strides.size() * sizeof(size_t), 5);
+    compute_encoder->setBytes(&ndim, sizeof(int), 6);
+
+    // Each thread group is responsible for 1 output
+    size_t n_reads = REDUCE_N_READS;
+    size_t thread_group_size =
+        row_reduce_kernel->maxTotalThreadsPerThreadgroup();
+    thread_group_size =
+        std::min((reduction_size + n_reads - 1) / n_reads, thread_group_size);
+
+    // Align thread group size with simd_size
+    uint simd_size = row_reduce_kernel->threadExecutionWidth();
+    thread_group_size =
+        (thread_group_size + simd_size - 1) / simd_size * simd_size;
+    assert(thread_group_size <= kernel->maxTotalThreadsPerThreadgroup());
+
+    // Launch enough thread groups for each output
+    uint n_threads = thread_group_size;
+    grid_dims = MTL::Size(n_threads, out.size(), 1);
+    group_dims = MTL::Size(thread_group_size, 1, 1);
+
+    compute_encoder->dispatchThreads(grid_dims, group_dims);
+
+    d.get_command_buffer(s.index)->addCompletedHandler(
+        [intermediates](MTL::CommandBuffer*) mutable {
+          intermediates.clear();
+        });
+  }
 }
 
 } // namespace
@@ -223,14 +460,6 @@ void Reduce::eval_gpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   array in = inputs[0];
 
-  // TODO: Allow specific row and column reductions with types disabled
-  // due to atomics ?
-  if (size_of(in.dtype()) == 8) {
-    std::ostringstream msg;
-    msg << "[Reduce::eval_gpu] Does not support " << in.dtype();
-    throw std::runtime_error(msg.str());
-  }
-
   // Make sure no identity reductions trickle down here
   assert(!axes_.empty());
 
@@ -297,7 +526,7 @@ void Reduce::eval_gpu(const std::vector<array>& inputs, array& out) {
     // Reducing over everything and the data is all there no broadcasting or
     // slicing etc.
     if (plan.type == ContiguousAllReduce) {
-      all_reduce_dispatch(in, out, op_name, compute_encoder, d);
+      all_reduce_dispatch(in, out, op_name, compute_encoder, d, s);
     }
 
     // At least the last dimension is row contiguous and we are reducing over
@@ -305,7 +534,7 @@ void Reduce::eval_gpu(const std::vector<array>& inputs, array& out) {
     else if (
         plan.type == ContiguousReduce || plan.type == GeneralContiguousReduce) {
       row_reduce_general_dispatch(
-          in, out, op_name, plan, axes_, compute_encoder, d);
+          in, out, op_name, plan, axes_, compute_encoder, d, s);
     }
 
     // At least the last two dimensions are contiguous and we are doing a
@@ -314,7 +543,7 @@ void Reduce::eval_gpu(const std::vector<array>& inputs, array& out) {
         plan.type == ContiguousStridedReduce ||
         plan.type == GeneralStridedReduce) {
       strided_reduce_general_dispatch(
-          in, out, op_name, plan, axes_, compute_encoder, d);
+          in, out, op_name, plan, axes_, compute_encoder, d, s);
     }
 
     if (!copies.empty()) {

mlx/backend/no_metal/primitives.cpp

@@ -1,4 +1,4 @@
-// Copyright © 2023 Apple Inc.
+// Copyright © 2023-2024 Apple Inc.
 
 #include "mlx/primitives.h"
 
@@ -37,6 +37,8 @@ NO_GPU(Convolution)
 NO_GPU(Copy)
 NO_GPU(Cos)
 NO_GPU(Cosh)
+NO_GPU_MULTI(CustomVJP)
+NO_GPU_MULTI(Depends)
 NO_GPU(Divide)
 NO_GPU(Remainder)
 NO_GPU(Equal)
@@ -90,5 +92,5 @@ NO_GPU(Tan)
 NO_GPU(Tanh)
 NO_GPU(Transpose)
 NO_GPU_MULTI(DivMod)
-
+NO_GPU_MULTI(QRF)
 } // namespace mlx::core

mlx/compile.cpp

@@ -0,0 +1,440 @@
+// Copyright © 2023-2024 Apple Inc.
+#include <cstdlib>
+#include <map>
+#include <unordered_map>
+#include <unordered_set>
+
+#include "mlx/allocator.h"
+#include "mlx/primitives.h"
+#include "mlx/transforms.h"
+#include "mlx/transforms_impl.h"
+
+namespace mlx::core {
+
+namespace detail {
+
+bool& compiler_disabled() {
+  auto get_val = []() {
+    if (const char* buff_str = std::getenv("MLX_DISABLE_COMPILE")) {
+      return true;
+    } else {
+      return false;
+    }
+  };
+  static bool compiler_disabled_ = get_val();
+  return compiler_disabled_;
+}
+
+#define MAX_OPS_PER_BUFFER max_ops_per_buffer()
+
+using CompileFn = std::function<std::vector<array>(const std::vector<array>&)>;
+using ParentsMap =
+    std::unordered_map<std::uintptr_t, std::vector<std::pair<array, int>>>;
+
+template <typename T, typename... U>
+size_t getAddress(std::function<T(U...)> f) {
+  typedef T(fnType)(U...);
+  fnType** fnPointer = f.template target<fnType*>();
+  if (fnPointer == nullptr) {
+    throw std::invalid_argument(
+        "[compile] Cannot compile a non-addressable function.");
+  }
+  return (size_t)*fnPointer;
+}
+
+struct CompilerCache {
+  struct CacheEntry {
+    std::vector<array> inputs;
+    std::vector<array> outputs;
+    std::vector<array> tape;
+    bool empty{true};
+  };
+
+  // Returns a reference to a CacheEntry which can be updated
+  // by the caller to avoid copying large tapes / inputs / outputs
+  CacheEntry& find(size_t fun_id, const std::vector<array>& inputs) {
+    // Try to find the entry
+    auto [entry_it, inserted] = cache_.insert({fun_id, {}});
+    auto& entries = entry_it->second;
+    auto is_match = [](const std::vector<array>& in1,
+                       const std::vector<array>& in2) {
+      if (in1.size() != in2.size()) {
+        throw std::runtime_error(
+            "[compiler] Got different number of inputs to function,"
+            " this should never happen.");
+      }
+      for (int i = 0; i < in1.size(); ++i) {
+        if (in1[i].shape() != in2[i].shape()) {
+          return false;
+        }
+        if (in1[i].dtype() != in2[i].dtype()) {
+          return false;
+        }
+      }
+      return true;
+    };
+
+    // Loop over entries and check inputs match i.e. shapes and types must be
+    // equal. Note this could get really slow if one compiles the same
+    // function with many different shapes. May want to store entries in a
+    // more easily searchable structure.
+    for (auto& entry : entries) {
+      // Check the inputs match and return if so
+      if (is_match(inputs, entry.inputs)) {
+        return entry;
+      }
+    }
+    // Otherwise append a new cache entry
+    entries.push_back(CacheEntry{});
+    return entries.back();
+  };
+
+  void erase(size_t fun_id) {
+    cache_.erase(fun_id);
+  }
+
+ private:
+  CompilerCache() {
+    // Make sure the allocator is fully
+    // initialized before the compiler cache
+    allocator::allocator();
+  }
+  friend CompilerCache& compiler_cache();
+  std::unordered_map<size_t, std::vector<CacheEntry>> cache_;
+};
+
+CompilerCache& compiler_cache() {
+  static CompilerCache compiler_cache_;
+  return compiler_cache_;
+}
+
+std::pair<std::vector<array>, std::vector<array>> compile_trace(
+    const std::function<std::vector<array>(const std::vector<array>&)>& fun,
+    const std::vector<array>& inputs) {
+  // Set the global tracing flag.
+  detail::InTracing in_tracing;
+
+  // Run the function on placeholder inputs
+  // to get compute graph
+  std::vector<array> tracer_inputs;
+  for (int i = 0; i < inputs.size(); ++i) {
+    array in(inputs[i].shape(), inputs[i].dtype(), nullptr, {});
+    in.set_tracer(true);
+    tracer_inputs.push_back(std::move(in));
+  }
+  return {tracer_inputs, fun(tracer_inputs)};
+}
+
+// Traverses the graph to build a tape and a map of array ids to their parents
+std::pair<std::vector<array>, ParentsMap> compile_dfs(
+    const std::vector<array>& inputs,
+    const std::vector<array>& outputs) {
+  std::function<void(const array&)> recurse;
+  std::vector<array> tape;
+  std::unordered_set<std::uintptr_t> input_set;
+  std::unordered_map<std::uintptr_t, std::vector<std::pair<array, int>>>
+      parents_map;
+  for (int i = 0; i < inputs.size(); ++i) {
+    auto in = inputs[i];
+    input_set.insert(in.id());
+  }
+
+  // DFS the graph to build the tape, and log parents and scalars
+  std::unordered_set<std::uintptr_t> cache;
+  recurse = [&](const array& a) {
+    auto id = a.id();
+    if (cache.find(id) != cache.end()) {
+      return;
+    }
+    for (int i = 0; i < a.inputs().size(); i++) {
+      auto& in = a.inputs()[i];
+      parents_map[in.id()].push_back({a, i});
+      for (auto& s : a.siblings()) {
+        parents_map[in.id()].push_back({s, i});
+      }
+      // Don't recurse on inputs (but add them to the tape for the purpose
+      // of future optimizations)
+      if (input_set.find(a.id()) == input_set.end()) {
+        recurse(in);
+      }
+    }
+    cache.insert(id);
+    for (auto& s : a.siblings()) {
+      cache.insert(s.id());
+    }
+    tape.push_back(a);
+  };
+  for (auto& a : outputs) {
+    recurse(a);
+  }
+  return {tape, parents_map};
+}
+
+// Simplify the tape. Note, this function modifies in-place both the tape and
+// the parents map to remove orphaned arrays
+void compile_simplify(
+    std::vector<array>& tape,
+    ParentsMap& parents_map,
+    const std::vector<array>& outputs,
+    int passes) {
+  // Helpers to identify identical scalars
+  std::map<std::pair<uint64_t, Dtype::Val>, array> scalars;
+  auto is_scalar = [](const array& a) {
+    return a.is_evaled() && a.ndim() == 0;
+  };
+  auto get_scalar_rep = [](const array& a) {
+    uint64_t v = 0;
+    int dtype;
+    switch (a.dtype().size) {
+      case 1:
+        v = *a.data<uint8_t>();
+        break;
+      case 4:
+        v = *a.data<uint32_t>();
+        break;
+      case 8:
+        v = *a.data<uint64_t>();
+        break;
+    }
+    return std::make_pair(v, a.dtype().val);
+  };
+
+  for (auto& a : tape) {
+    if (is_scalar(a)) {
+      scalars.insert({get_scalar_rep(a), a});
+    }
+  }
+
+  // Helper that fuses two arrays in the graph by setting the parents of the
+  // source to point to the destination
+  auto fuse = [&](array& dst, array& src) {
+    // Canonicalize the order of the primitives outputs
+    auto sources = src.outputs();
+    auto dests = dst.outputs();
+    // For each src parent, point it to the corresponding dest
+    for (int i = 0; i < sources.size(); ++i) {
+      auto src_parents = parents_map.find(sources[i].id());
+      if (src_parents == parents_map.end()) {
+        continue;
+      }
+      auto& pairs = parents_map[dests[i].id()];
+      for (auto& parent : src_parents->second) {
+        parent.first.inputs()[parent.second] = dests[i];
+        pairs.push_back(parent);
+      }
+      // Remove the source from the map to avoid fusing with it again
+      parents_map.erase(src_parents);
+    }
+  };
+
+  // Depth-1 array equivalence check.
+  auto array_equivalent = [](const array& a, const array& b) {
+    if (!a.has_primitive() || !b.has_primitive()) {
+      return false;
+    }
+    if (a.primitive_id() == b.primitive_id()) {
+      return false;
+    }
+    const auto& pa = a.primitive();
+    const auto& pb = b.primitive();
+    if (typeid(pa) != typeid(pb)) {
+      return false;
+    }
+
+    if (a.inputs().size() != b.inputs().size()) {
+      return false;
+    }
+
+    for (int i = 0; i < a.inputs().size(); i++) {
+      if (a.inputs()[i].id() != b.inputs()[i].id()) {
+        return false;
+      }
+    }
+
+    return pa.is_equivalent(pb);
+  };
+
+  // Pass 0: fuse scalars
+  std::vector<array> new_tape;
+  for (auto& arr : tape) {
+    // Check if we can fuse scalars
+    if (is_scalar(arr)) {
+      auto scalar = scalars.find(get_scalar_rep(arr));
+      if (scalar->second.id() != arr.id()) {
+        fuse(scalar->second, arr);
+        // Don't keep orphaned scalars in the tape
+        continue;
+      }
+    }
+    new_tape.push_back(std::move(arr));
+  }
+
+  tape = std::move(new_tape);
+
+  std::unordered_set<uintptr_t> output_set;
+  for (auto& o : outputs) {
+    output_set.insert(o.id());
+  }
+  // Pass 1..passes: fuse only keeping non-orphaned arrays in the tape
+  for (int pass = 0; pass < passes; ++pass) {
+    for (auto& arr : tape) {
+      // Helper to check if we can fuse the parents of the
+      // given array
+      auto maybe_fuse_parents = [&](auto& a) {
+        auto parents = parents_map.find(a.id());
+        if (parents != parents_map.end()) {
+          auto N = parents->second.size();
+          std::vector<bool> mask(N, false);
+          for (int i = 0; i < N; i++) {
+            if (mask[i]) {
+              continue;
+            }
+            for (int j = i + 1; j < N; j++) {
+              if (mask[j]) {
+                continue;
+              }
+              auto& src = parents->second[j].first;
+              auto& dst = parents->second[i].first;
+              if (src.id() != dst.id() && array_equivalent(src, dst)) {
+                fuse(dst, src);
+                mask[j] = true;
+              }
+            }
+          }
+          // Erase orphaned parents so we don't keep fusing with them
+          for (int i = N - 1; i > 0; --i) {
+            if (mask[i]) {
+              parents->second.erase(parents->second.begin() + i);
+            }
+          }
+          return false;
+        } else {
+          return output_set.find(a.id()) == output_set.end();
+        }
+      };
+
+      bool discard = maybe_fuse_parents(arr);
+      for (auto& s : arr.siblings()) {
+        discard &= maybe_fuse_parents(s);
+      }
+      // If an array and its siblings have no parents, and none of them are
+      // outputs, it is safe to remove it from the tape
+      if (!discard) {
+        new_tape.push_back(std::move(arr));
+      }
+    }
+    tape = std::move(new_tape);
+  }
+}
+
+std::vector<array> compile_replace(
+    const std::vector<array>& tape,
+    const std::vector<array>& trace_inputs,
+    const std::vector<array>& trace_outputs,
+    const std::vector<array>& inputs) {
+  std::unordered_map<uintptr_t, array> trace_to_real;
+  for (int i = 0; i < inputs.size(); ++i) {
+    trace_to_real.insert({trace_inputs[i].id(), inputs[i]});
+  }
+
+  for (auto& a : tape) {
+    // Arrays in the tape without primitives are constants
+    // and can be used directly
+    if (!a.has_primitive()) {
+      trace_to_real.insert({a.id(), a});
+    } else {
+      // Find real inputs
+      std::vector<array> real_inputs;
+      for (auto& in : a.inputs()) {
+        real_inputs.push_back(trace_to_real.at(in.id()));
+      }
+      if (a.siblings().empty()) {
+        auto real_a = array(
+            a.shape(), a.dtype(), a.primitive_ptr(), std::move(real_inputs));
+        trace_to_real.insert({a.id(), std::move(real_a)});
+      } else {
+        // Ensure the order is correct for multi-output primitives
+        std::vector<std::vector<int>> shapes;
+        std::vector<Dtype> types;
+        auto trace_out = a.outputs();
+        for (auto& o : trace_out) {
+          shapes.push_back(o.shape());
+          types.push_back(o.dtype());
+        }
+        auto real_out =
+            array::make_arrays(shapes, types, a.primitive_ptr(), real_inputs);
+        for (int i = 0; i < trace_out.size(); ++i) {
+          trace_to_real.insert({trace_out[i].id(), std::move(real_out[i])});
+        }
+      }
+    }
+  }
+
+  std::vector<array> outputs;
+  for (auto& o : trace_outputs) {
+    outputs.push_back(trace_to_real.at(o.id()));
+  }
+  return outputs;
+}
+
+std::function<std::vector<array>(const std::vector<array>&)> compile(
+    const std::function<std::vector<array>(const std::vector<array>&)>& fun,
+    size_t fun_id) {
+  if (compiler_disabled()) {
+    return fun;
+  }
+  return [fun, fun_id](const std::vector<array>& inputs) {
+    // Find a cache entry with the correct inputs
+    auto& entry = compiler_cache().find(fun_id, inputs);
+
+    // No matching cache entry existed, so compile
+    if (entry.empty) {
+      // Mark the entry as not empty since we are about to fill it
+      entry.empty = false;
+      // Trace to build the graph
+      std::tie(entry.inputs, entry.outputs) = compile_trace(fun, inputs);
+
+      // DFS the graph and get a tape, and a map of array id to (parent,
+      // position in parent inputs)
+      std::unordered_map<uintptr_t, std::vector<std::pair<array, int>>>
+          parents_map;
+      std::tie(entry.tape, parents_map) =
+          compile_dfs(entry.inputs, entry.outputs);
+
+      // Simplify the tape
+      compile_simplify(entry.tape, parents_map, entry.outputs, /* passes */ 3);
+
+      // This is a good point to do more optimizations, e.g. kernel fusion to
+      // generate new primitives. The tape needs to be updated accordingly
+    }
+
+    // At this point we must have a tape, now replace the placeholders
+    // with real arrays that can be evaluated
+    return compile_replace(entry.tape, entry.inputs, entry.outputs, inputs);
+  };
+}
+
+void compile_erase(size_t fun_id) {
+  detail::compiler_cache().erase(fun_id);
+}
+
+} // namespace detail
+
+std::function<std::vector<array>(const std::vector<array>&)> compile(
+    const std::function<std::vector<array>(const std::vector<array>&)>& fun) {
+  if (detail::compiler_disabled()) {
+    return fun;
+  }
+  auto fun_id = detail::getAddress(fun);
+  return detail::compile(fun, fun_id);
+}
+
+void disable_compile() {
+  detail::compiler_disabled() = true;
+}
+
+void enable_compile() {
+  detail::compiler_disabled() = false;
+}
+
+} // namespace mlx::core

mlx/io.h

@@ -0,0 +1,55 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#pragma once
+
+#include <variant>
+
+#include "mlx/array.h"
+#include "mlx/io/load.h"
+#include "mlx/ops.h"
+#include "mlx/stream.h"
+
+namespace mlx::core {
+
+/** Save array to out stream in .npy format */
+void save(std::shared_ptr<io::Writer> out_stream, array a);
+
+/** Save array to file in .npy format */
+void save(const std::string& file, array a);
+
+/** Load array from reader in .npy format */
+array load(std::shared_ptr<io::Reader> in_stream, StreamOrDevice s = {});
+
+/** Load array from file in .npy format */
+array load(const std::string& file, StreamOrDevice s = {});
+
+/** Load array map from .safetensors file format */
+std::unordered_map<std::string, array> load_safetensors(
+    std::shared_ptr<io::Reader> in_stream,
+    StreamOrDevice s = {});
+std::unordered_map<std::string, array> load_safetensors(
+    const std::string& file,
+    StreamOrDevice s = {});
+
+void save_safetensors(
+    std::shared_ptr<io::Writer> in_stream,
+    std::unordered_map<std::string, array>);
+void save_safetensors(
+    const std::string& file,
+    std::unordered_map<std::string, array>);
+
+using MetaData =
+    std::variant<std::monostate, array, std::string, std::vector<std::string>>;
+
+/** Load array map and metadata from .gguf file format */
+std::pair<
+    std::unordered_map<std::string, array>,
+    std::unordered_map<std::string, MetaData>>
+load_gguf(const std::string& file, StreamOrDevice s = {});
+
+void save_gguf(
+    std::string file,
+    std::unordered_map<std::string, array> array_map,
+    std::unordered_map<std::string, MetaData> meta_data = {});
+
+} // namespace mlx::core

mlx/io/CMakeLists.txt

@@ -4,6 +4,7 @@ target_sources(
   ${CMAKE_CURRENT_SOURCE_DIR}/load.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/safetensor.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/gguf.cpp
+  ${CMAKE_CURRENT_SOURCE_DIR}/gguf_quants.cpp
 )
 
 MESSAGE(STATUS "Downloading json")

mlx/io/gguf.cpp

@@ -1,17 +1,16 @@
-// Copyright © 2023 Apple Inc.
+// Copyright © 2023-2024 Apple Inc.
 
+#include <cstdint>
 #include <cstring>
+#include <numeric>
 
-#include "mlx/ops.h"
-#include "mlx/primitives.h"
-#include "mlx/utils.h"
-
-extern "C" {
-#include <gguflib.h>
-}
+#include <mlx/io/gguf.h>
 
 namespace mlx::core {
 
+// https://github.com/antirez/gguf-tools/blob/af7d88d808a7608a33723fba067036202910acb3/gguflib.h#L102-L108
+constexpr int gguf_array_header_size = 12;
+
 std::optional<uint32_t> dtype_to_gguf_tensor_type(const Dtype& dtype) {
   switch (dtype) {
     case float32:
@@ -46,6 +45,15 @@ std::optional<Dtype> gguf_type_to_dtype(const uint32_t& gguf_type) {
   }
 }
 
+std::vector<int> get_shape(const gguf_tensor& tensor) {
+  std::vector<int> shape;
+  // The dimension order in GGML is the reverse of the order used in MLX.
+  for (int i = tensor.ndim - 1; i >= 0; i--) {
+    shape.push_back(tensor.dim[i]);
+  }
+  return shape;
+}
+
 std::tuple<allocator::Buffer, Dtype> extract_tensor_data(gguf_tensor* tensor) {
   std::optional<Dtype> equivalent_dtype = gguf_type_to_dtype(tensor->type);
   // If there's an equivalent type, we can simply copy.
@@ -70,46 +78,328 @@ std::tuple<allocator::Buffer, Dtype> extract_tensor_data(gguf_tensor* tensor) {
   return {buffer, float16};
 }
 
-std::unordered_map<std::string, array> load_gguf(
-    const std::string& file,
-    StreamOrDevice s) {
-  std::unordered_map<std::string, array> result;
+void set_mx_value_from_gguf(
+    gguf_ctx* ctx,
+    uint32_t type,
+    gguf_value* val,
+    MetaData& value) {
+  switch (type) {
+    case GGUF_VALUE_TYPE_UINT8:
+      value = array(val->uint8, uint8);
+      break;
+    case GGUF_VALUE_TYPE_INT8:
+      value = array(val->int8, int8);
+      break;
+    case GGUF_VALUE_TYPE_UINT16:
+      value = array(val->uint16, uint16);
+      break;
+    case GGUF_VALUE_TYPE_INT16:
+      value = array(val->int16, int16);
+      break;
+    case GGUF_VALUE_TYPE_UINT32:
+      value = array(val->uint32, uint32);
+      break;
+    case GGUF_VALUE_TYPE_INT32:
+      value = array(val->int32, int32);
+      break;
+    case GGUF_VALUE_TYPE_UINT64:
+      value = array(val->uint64, uint64);
+      break;
+    case GGUF_VALUE_TYPE_INT64:
+      value = array(val->int64, int64);
+      break;
+    case GGUF_VALUE_TYPE_FLOAT32:
+      value = array(val->float32, float32);
+      break;
+    case GGUF_VALUE_TYPE_BOOL:
+      value = array(val->boolval, bool_);
+      break;
+    case GGUF_VALUE_TYPE_STRING:
+      value =
+          std::string(val->string.string, static_cast<int>(val->string.len));
+      break;
+    case GGUF_VALUE_TYPE_FLOAT64:
+      value = array(val->float64, float32);
+      break;
+    case GGUF_VALUE_TYPE_ARRAY: {
+      ctx->off += gguf_array_header_size; // Skip header
+      char* data = reinterpret_cast<char*>(val) + gguf_array_header_size;
+      auto size = static_cast<int>(val->array.len);
+      if (val->array.type == GGUF_VALUE_TYPE_ARRAY) {
+        throw std::invalid_argument(
+            "[load_gguf] Only supports loading 1-layer of nested arrays.");
+      }
+      switch (val->array.type) {
+        case GGUF_VALUE_TYPE_UINT8:
+          value = array(reinterpret_cast<uint8_t*>(data), {size}, uint8);
+          break;
+        case GGUF_VALUE_TYPE_INT8:
+          value = array(reinterpret_cast<int8_t*>(data), {size}, int8);
+          break;
+        case GGUF_VALUE_TYPE_UINT16:
+          value = array(reinterpret_cast<uint16_t*>(data), {size}, uint16);
+          break;
+        case GGUF_VALUE_TYPE_INT16:
+          value = array(reinterpret_cast<int16_t*>(data), {size}, int16);
+          break;
+        case GGUF_VALUE_TYPE_UINT32:
+          value = array(reinterpret_cast<uint32_t*>(data), {size}, uint32);
+          break;
+        case GGUF_VALUE_TYPE_INT32:
+          value = array(reinterpret_cast<int32_t*>(data), {size}, int32);
+          break;
+        case GGUF_VALUE_TYPE_UINT64:
+          value = array(reinterpret_cast<uint64_t*>(data), {size}, uint64);
+          break;
+        case GGUF_VALUE_TYPE_INT64:
+          value = array(reinterpret_cast<uint64_t*>(data), {size}, int64);
+          break;
+        case GGUF_VALUE_TYPE_FLOAT32:
+          value = array(reinterpret_cast<float*>(data), {size}, float32);
+          break;
+        case GGUF_VALUE_TYPE_BOOL:
+          value = array(reinterpret_cast<bool*>(data), {size}, bool_);
+          break;
+        case GGUF_VALUE_TYPE_STRING: {
+          std::vector<std::string> strs(size);
+          for (auto& str : strs) {
+            auto str_val = reinterpret_cast<gguf_string*>(data);
+            data += (str_val->len + sizeof(gguf_string));
+            str = std::string(str_val->string, static_cast<int>(str_val->len));
+            ctx->off += (str_val->len + sizeof(gguf_string));
+          }
+          value = std::move(strs);
+          break;
+        }
+        case GGUF_VALUE_TYPE_FLOAT64:
+          value = array(reinterpret_cast<double*>(data), {size}, float32);
+          break;
+        default:
+          throw std::runtime_error(
+              "[load_gguf] Multiple levels of nested arrays are not supported.");
+      }
+      break;
+    }
+    default:
+      throw std::runtime_error("[load_gguf] Received unexpected type.");
+      break;
+  }
+  if (type == GGUF_VALUE_TYPE_STRING) {
+    ctx->off += (sizeof(gguf_string) + std::get<std::string>(value).size());
+  } else if (auto pv = std::get_if<array>(&value); pv) {
+    ctx->off += pv->nbytes();
+  }
+}
+
+std::unordered_map<std::string, MetaData> load_metadata(gguf_ctx* ctx) {
+  std::unordered_map<std::string, MetaData> metadata;
+  gguf_key key;
+  while (gguf_get_key(ctx, &key)) {
+    std::string key_name = std::string(key.name, key.namelen);
+    auto& val = metadata.insert({key_name, MetaData{}}).first->second;
+    set_mx_value_from_gguf(ctx, key.type, key.val, val);
+  }
+  return metadata;
+}
+
+std::unordered_map<std::string, array> load_arrays(gguf_ctx* ctx) {
+  std::unordered_map<std::string, array> array_map;
+  gguf_tensor tensor;
+
+  auto check_insert = [](auto inserted) {
+    if (!inserted.second) {
+      std::ostringstream msg;
+      msg << "[load_gguf] Duplicate parameter name " << inserted.first->second
+          << " this can happend when loading quantized tensors.";
+      throw std::runtime_error(msg.str());
+    }
+  };
+
+  while (gguf_get_tensor(ctx, &tensor)) {
+    if (tensor.type == GGUF_TYPE_Q4_0 || tensor.type == GGUF_TYPE_Q4_1 ||
+        tensor.type == GGUF_TYPE_Q8_0) {
+      gguf_load_quantized(array_map, tensor);
+    } else {
+      std::string name = std::string(tensor.name, tensor.namelen);
+
+      const auto& [data, dtype] = extract_tensor_data(&tensor);
+      array loaded_array = array(data, get_shape(tensor), dtype);
+      array_map.insert({name, loaded_array});
+    }
+  }
+  return array_map;
+}
+
+std::pair<
+    std::unordered_map<std::string, array>,
+    std::unordered_map<std::string, MetaData>>
+load_gguf(const std::string& file, StreamOrDevice s) {
   gguf_ctx* ctx = gguf_open(file.c_str());
   if (!ctx) {
     throw std::runtime_error("[load_gguf] gguf_init failed");
   }
-  gguf_skip_key_values_section(ctx);
-  gguf_tensor tensor;
-  while (gguf_get_tensor(ctx, &tensor)) {
-    std::vector<int> shape;
-    // The dimension order in GGML is the reverse of the order used in MLX.
-    for (int i = tensor.ndim - 1; i >= 0; i--) {
-      shape.push_back(tensor.dim[i]);
-    }
-    const auto& [data, dtype] = extract_tensor_data(&tensor);
-    array loaded_array = array(data, shape, dtype);
-    std::string name = std::string(tensor.name, tensor.namelen);
-    result.insert({name, loaded_array});
-  }
+  auto metadata = load_metadata(ctx);
+  auto arrays = load_arrays(ctx);
   gguf_close(ctx);
-  return result;
+  return {arrays, metadata};
 }
 
-void save_gguf(std::string file, std::unordered_map<std::string, array> a) {
+void append_kv_array(
+    gguf_ctx* ctx,
+    const std::string& key,
+    array& val,
+    uint32_t gguf_type) {
+  if (val.ndim() == 1) {
+    size_t gguf_size = val.nbytes() + gguf_array_header_size;
+    std::vector<char> val_vec(gguf_size);
+    gguf_value* gguf_val = reinterpret_cast<gguf_value*>(val_vec.data());
+    gguf_val->array.type = gguf_type;
+    gguf_val->array.len = val.size();
+    memcpy(
+        val_vec.data() + gguf_array_header_size,
+        val.data<char>(),
+        val.nbytes());
+    gguf_append_kv(
+        ctx,
+        key.c_str(),
+        key.length(),
+        GGUF_VALUE_TYPE_ARRAY,
+        reinterpret_cast<void*>(val_vec.data()),
+        gguf_size);
+  } else {
+    gguf_append_kv(
+        ctx,
+        key.c_str(),
+        key.length(),
+        gguf_type,
+        reinterpret_cast<void*>(val.data<char>()),
+        val.nbytes());
+  }
+}
+
+void save_gguf(
+    std::string file,
+    std::unordered_map<std::string, array> array_map,
+    std::unordered_map<std::string, MetaData> metadata /* = {} */) {
   // Add .gguf to file name if it is not there
   if (file.length() < 5 || file.substr(file.length() - 5, 5) != ".gguf") {
     file += ".gguf";
   }
+
   gguf_ctx* ctx = gguf_create(file.c_str(), GGUF_OVERWRITE);
   if (!ctx) {
     throw std::runtime_error("[save_gguf] gguf_create failed");
   }
 
+  auto string_to_gguf = [](char* dst, const std::string& src) {
+    gguf_string* val = reinterpret_cast<gguf_string*>(dst);
+    val->len = src.length();
+    memcpy(val->string, src.c_str(), src.length());
+  };
+
+  // Save any meta data
+  for (auto& [key, value] : metadata) {
+    if (auto pv = std::get_if<std::string>(&value); pv) {
+      const std::string& str = *pv;
+      size_t size = sizeof(gguf_string) + str.length();
+      std::vector<char> val_vec(size);
+      string_to_gguf(val_vec.data(), str);
+      gguf_append_kv(
+          ctx,
+          key.c_str(),
+          key.length(),
+          GGUF_VALUE_TYPE_STRING,
+          static_cast<void*>(val_vec.data()),
+          size);
+    } else if (auto pv = std::get_if<std::vector<std::string>>(&value); pv) {
+      const auto& str_vec = *pv;
+      auto mem_size = std::accumulate(
+          str_vec.begin(), str_vec.end(), 0, [](size_t accum, const auto& s) {
+            return accum + s.size();
+          });
+      mem_size += str_vec.size() * sizeof(gguf_string) + gguf_array_header_size;
+      std::vector<char> val_vec(mem_size);
+      gguf_value* val = reinterpret_cast<gguf_value*>(val_vec.data());
+      val->array.type = GGUF_VALUE_TYPE_STRING;
+      val->array.len = str_vec.size();
+      auto str_ptr = val_vec.data() + gguf_array_header_size;
+      for (auto& str : str_vec) {
+        string_to_gguf(str_ptr, str);
+        str_ptr += str.length() + sizeof(gguf_string);
+      }
+      gguf_append_kv(
+          ctx,
+          key.c_str(),
+          key.length(),
+          GGUF_VALUE_TYPE_ARRAY,
+          static_cast<void*>(val),
+          mem_size);
+    } else if (auto pv = std::get_if<array>(&value); pv) {
+      array v = *pv;
+      if (v.ndim() > 1) {
+        throw std::runtime_error(
+            "[save_gguf] Cannot save arrays with more than one dimension.");
+      }
+      if (v.size() == 0) {
+        throw std::runtime_error("[save_gguf] Cannot save empty arrays.");
+      }
+
+      eval(v);
+      if (!v.flags().row_contiguous) {
+        v = reshape(flatten(v), v.shape());
+      }
+      if (!v.flags().row_contiguous) {
+        throw std::runtime_error(
+            "[save_gguf] Cannot save non contiguous arrays.");
+      }
+      switch (v.dtype()) {
+        case float32:
+          append_kv_array(ctx, key, v, GGUF_VALUE_TYPE_FLOAT32);
+          break;
+        case int64:
+          append_kv_array(ctx, key, v, GGUF_VALUE_TYPE_INT64);
+          break;
+        case int32:
+          append_kv_array(ctx, key, v, GGUF_VALUE_TYPE_INT32);
+          break;
+        case int16:
+          append_kv_array(ctx, key, v, GGUF_VALUE_TYPE_INT16);
+          break;
+        case int8:
+          append_kv_array(ctx, key, v, GGUF_VALUE_TYPE_INT8);
+          break;
+        case uint64:
+          append_kv_array(ctx, key, v, GGUF_VALUE_TYPE_UINT64);
+          break;
+        case uint32:
+          append_kv_array(ctx, key, v, GGUF_VALUE_TYPE_UINT32);
+          break;
+        case uint16:
+          append_kv_array(ctx, key, v, GGUF_VALUE_TYPE_UINT16);
+          break;
+        case uint8:
+          append_kv_array(ctx, key, v, GGUF_VALUE_TYPE_UINT8);
+          break;
+        case bool_:
+          append_kv_array(ctx, key, v, GGUF_VALUE_TYPE_BOOL);
+          break;
+        default:
+          std::ostringstream msg;
+          msg << "[save_gguf] array type " << v.dtype()
+              << " not support for metadata.";
+          throw std::invalid_argument(msg.str());
+      }
+    } else {
+      throw std::runtime_error(
+          "[save_gguf] Received unexpected type in metadata");
+    }
+  }
+
   // Tensor offsets are relative to data section, so we start at offset 0.
   uint64_t tensor_offset = 0;
 
   // First, append the tensor info
-  for (auto& [key, arr] : a) {
+  for (auto& [key, arr] : array_map) {
     arr.eval();
 
     // Try to make it row contiguous
@@ -154,7 +444,7 @@ void save_gguf(std::string file, std::unordered_map<std::string, array> a) {
   }
 
   // Then, append the tensor weights
-  for (const auto& [key, arr] : a) {
+  for (const auto& [key, arr] : array_map) {
     if (!gguf_append_tensor_data(ctx, (void*)arr.data<void>(), arr.nbytes())) {
       throw std::runtime_error("[save_gguf] gguf_append_tensor_data failed");
     }

mlx/io/gguf.h

@@ -0,0 +1,20 @@
+// Copyright © 2023-2024 Apple Inc.
+#pragma once
+
+#include "mlx/io.h"
+#include "mlx/primitives.h"
+#include "mlx/transforms.h"
+#include "mlx/utils.h"
+
+extern "C" {
+#include <gguflib.h>
+}
+
+namespace mlx::core {
+
+std::vector<int> get_shape(const gguf_tensor& tensor);
+void gguf_load_quantized(
+    std::unordered_map<std::string, array>& a,
+    const gguf_tensor& tensor);
+
+} // namespace mlx::core

mlx/io/gguf_quants.cpp

@@ -0,0 +1,158 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include <cstdint>
+#include <cstring>
+
+#include <mlx/io/gguf.h>
+
+namespace mlx::core {
+
+void unpack_32_4(uint8_t* data, int8_t* dst) {
+  for (int64_t j = 0; j < 16; ++j) {
+    uint8_t x = (data[j + 2] & 0x0F); // j+2 to skip scale bytes.
+    if (j % 2 != 0) {
+      x <<= 4;
+    }
+    dst[j / 2] += x;
+  }
+  // Last 16 weights are in the higher bits
+  for (int64_t j = 0; j < 16; ++j) {
+    uint8_t x = (data[j + 2] >> 4);
+    if (j % 2 != 0) {
+      x <<= 4;
+    }
+    dst[8 + j / 2] += x;
+  }
+}
+
+// Extracts (weight, scales, biases) from Q4_0 tensors.
+// Data layout is: |16 bit scale|32 x 4bit weights|.
+void extract_q4_0_data(
+    const gguf_tensor& tensor,
+    array& weights_arr,
+    array& scales_arr,
+    array& biases_arr) {
+  const uint64_t bytes_per_block = 18; // 2 bytes scale, 32x0.5 byte weights
+  auto data = static_cast<uint8_t*>(tensor.weights_data);
+  auto weights = weights_arr.data<int8_t>();
+  auto scales = scales_arr.data<float16_t>();
+  auto biases = biases_arr.data<float16_t>();
+  for (int64_t i = 0; i < scales_arr.size(); i++) {
+    scales[i] = *((float16_t*)data);
+    biases[i] = -8 * scales[i];
+    unpack_32_4(data, weights);
+    weights += 16;
+    data += bytes_per_block;
+  }
+}
+
+// Extracts (weight, scales, biases) from Q4_1 tensors.
+// Data layout is: |16 bit scale|16 bit bias|32 x 4bit weights|.
+void extract_q4_1_data(
+    const gguf_tensor& tensor,
+    array& weights_arr,
+    array& scales_arr,
+    array& biases_arr) {
+  const uint64_t bytes_per_block =
+      20; // 2 bytes scale, 2 bytes bias, 32x0.5 byte weights
+  auto data = static_cast<uint8_t*>(tensor.weights_data);
+  auto weights = weights_arr.data<int8_t>();
+  auto scales = scales_arr.data<float16_t>();
+  auto biases = biases_arr.data<float16_t>();
+  for (int64_t i = 0; i < scales_arr.size(); i++) {
+    scales[i] = *((float16_t*)data);
+    biases[i] = *((float16_t*)(data) + 1);
+    unpack_32_4(data, weights);
+    weights += 16;
+    data += bytes_per_block;
+  }
+}
+
+// Extracts (weight, scales, biases) from Q8_0 tensors.
+// Data layout is: |16 bit scale|32 x 8bit weights|.
+void extract_q8_0_data(
+    const gguf_tensor& tensor,
+    array& weights_arr,
+    array& scales_arr,
+    array& biases_arr) {
+  const uint64_t weights_per_block = 32;
+  const uint64_t bytes_per_block = 34; // 2 bytes scale, 32x1 byte weights
+  auto data = static_cast<uint8_t*>(tensor.weights_data);
+  auto weights = weights_arr.data<int8_t>();
+  auto scales = scales_arr.data<float16_t>();
+  auto biases = biases_arr.data<float16_t>();
+  for (int64_t i = 0; i < scales_arr.size(); i++) {
+    uint8_t* block_data = data + i * bytes_per_block;
+    scales[i] = *((float16_t*)block_data);
+    biases[i] = -128 * scales[i];
+    for (int64_t j = 0; j < weights_per_block; ++j) {
+      uint8_t x = block_data[j + 2]; // j+2 to skip the scale bytes.
+      // Original data is in int8_t, so we add a bias of -128 and invert the
+      // first bit.
+      x ^= 1 << 7;
+      weights[i * weights_per_block + j] = x;
+    }
+  }
+}
+
+void gguf_load_quantized(
+    std::unordered_map<std::string, array>& a,
+    const gguf_tensor& tensor) {
+  uint64_t weights_per_byte;
+  if (tensor.type == GGUF_TYPE_Q4_0 || tensor.type == GGUF_TYPE_Q4_1) {
+    weights_per_byte = 2;
+  } else { // tensor.type == GGUF_TYPE_Q8_0
+    weights_per_byte = 1;
+  }
+
+  std::string name = std::string(tensor.name, tensor.namelen);
+  std::vector<int> shape = get_shape(tensor);
+  const uint64_t weights_per_block = 32;
+  if (shape[shape.size() - 1] % weights_per_block != 0) {
+    std::ostringstream msg;
+    msg << "[load_gguf] tensor " << name
+        << "has incompatible last dim shape: " << shape[shape.size() - 1];
+    throw std::runtime_error(msg.str());
+  }
+  const uint64_t num_blocks = tensor.num_weights / weights_per_block;
+
+  std::vector<int> weights_shape = shape;
+  weights_shape.back() /= (weights_per_byte * 4);
+
+  array weights(std::move(weights_shape), uint32, nullptr, {});
+  weights.set_data(allocator::malloc(weights.nbytes()));
+
+  // For scales and bias
+  shape[shape.size() - 1] = shape[shape.size() - 1] / weights_per_block;
+  array scales(shape, float16, nullptr, {});
+  array biases(std::move(shape), float16, nullptr, {});
+  scales.set_data(allocator::malloc(scales.nbytes()));
+  biases.set_data(allocator::malloc(biases.nbytes()));
+
+  if (tensor.type == GGUF_TYPE_Q4_0) {
+    extract_q4_0_data(tensor, weights, scales, biases);
+  } else if (tensor.type == GGUF_TYPE_Q4_1) {
+    extract_q4_1_data(tensor, weights, scales, biases);
+  } else if (tensor.type == GGUF_TYPE_Q8_0) {
+    extract_q8_0_data(tensor, weights, scales, biases);
+  }
+
+  a.insert({name, weights});
+
+  auto check_insert = [](auto inserted) {
+    if (!inserted.second) {
+      std::ostringstream msg;
+      msg << "[load_gguf] Duplicate parameter name " << inserted.first->second
+          << " this can happend when loading quantized tensors.";
+      throw std::runtime_error(msg.str());
+    }
+  };
+
+  const std::string weight_suffix = ".weight";
+  const std::string name_prefix =
+      name.substr(0, name.length() - weight_suffix.length());
+  check_insert(a.insert({name_prefix + ".scales", scales}));
+  check_insert(a.insert({name_prefix + ".biases", biases}));
+}
+
+} // namespace mlx::core

mlx/io/safetensor.cpp

@@ -3,8 +3,8 @@
 #include <json.hpp>
 #include <stack>
 
+#include "mlx/io.h"
 #include "mlx/io/load.h"
-#include "mlx/ops.h"
 #include "mlx/primitives.h"
 
 using json = nlohmann::json;

mlx/linalg.cpp

@@ -4,8 +4,9 @@
 #include <ostream>
 #include <vector>
 
-#include "mlx/dtype.h"
 #include "mlx/linalg.h"
+#include "mlx/primitives.h"
+#include "mlx/utils.h"
 
 namespace mlx::core::linalg {
 
@@ -172,4 +173,31 @@ array norm(
   return matrix_norm(a, ord, ax, keepdims, s);
 }
 
+std::pair<array, array> qr(const array& a, StreamOrDevice s /* = {} */) {
+  if (a.dtype() != float32) {
+    std::ostringstream msg;
+    msg << "[linalg::qr] Arrays must type float32. Received array "
+        << "with type " << a.dtype() << ".";
+    throw std::invalid_argument(msg.str());
+  }
+  if (a.ndim() < 2) {
+    std::ostringstream msg;
+    msg << "[linalg::qr] 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(
+        "[linalg::qr] Support for non-square matrices NYI.");
+  }
+
+  auto out = array::make_arrays(
+      {a.shape(), a.shape()},
+      {a.dtype(), a.dtype()},
+      std::make_unique<QRF>(to_stream(s)),
+      {astype(a, a.dtype(), s)});
+  return std::make_pair(out[0], out[1]);
+}
+
 } // namespace mlx::core::linalg

mlx/linalg.h

@@ -60,4 +60,6 @@ norm(const array& a, int axis, bool keepdims = false, StreamOrDevice s = {}) {
   return norm(a, std::vector<int>{axis}, keepdims, s);
 }
 
+std::pair<array, array> qr(const array& a, StreamOrDevice s = {});
+
 } // namespace mlx::core::linalg

mlx/mlx.h

@@ -6,6 +6,7 @@
 #include "mlx/backend/metal/metal.h"
 #include "mlx/device.h"
 #include "mlx/fft.h"
+#include "mlx/io.h"
 #include "mlx/linalg.h"
 #include "mlx/ops.h"
 #include "mlx/random.h"

mlx/ops.cpp

@@ -1,4 +1,5 @@
-// Copyright © 2023 Apple Inc.
+// Copyright © 2023-2024 Apple Inc.
+
 #include <algorithm>
 #include <cmath>
 #include <numeric>
@@ -7,6 +8,7 @@
 
 #include "mlx/ops.h"
 #include "mlx/primitives.h"
+#include "mlx/transforms.h"
 #include "mlx/utils.h"
 
 namespace mlx::core {
@@ -15,8 +17,7 @@ namespace {
 
 std::pair<std::vector<int>, std::vector<int>> compute_reduce_shape(
     const std::vector<int>& axes,
-    const std::vector<int>& shape,
-    bool keepdims) {
+    const std::vector<int>& shape) {
   std::set<int> axes_set;
   auto ndim = shape.size();
   for (auto ax : axes) {
@@ -36,7 +37,7 @@ std::pair<std::vector<int>, std::vector<int>> compute_reduce_shape(
   for (int i = 0; i < ndim; ++i) {
     if (axes_set.count(i) == 0) {
       out_shape.push_back(shape[i]);
-    } else if (keepdims) {
+    } else {
       out_shape.push_back(1);
     }
   }
@@ -79,7 +80,14 @@ array arange(
     msg << bool_ << " not supported for arange.";
     throw std::invalid_argument(msg.str());
   }
-  int size = std::max(static_cast<int>(std::ceil((stop - start) / step)), 0);
+  if (std::isnan(start) || std::isnan(step) || std::isnan(stop)) {
+    throw std::invalid_argument("[arange] Cannot compute length.");
+  }
+  double real_size = std::ceil((stop - start) / step);
+  if (std::isnan(real_size)) {
+    throw std::invalid_argument("[arange] Cannot compute length.");
+  }
+  int size = std::max(static_cast<int>(real_size), 0);
   return array(
       {size},
       dtype,
@@ -182,6 +190,9 @@ array full(
     const array& vals,
     Dtype dtype,
     StreamOrDevice s /* = {} */) {
+  if (std::any_of(shape.begin(), shape.end(), [](auto i) { return i < 0; })) {
+    throw std::invalid_argument("[full] Negative dimensions not allowed.");
+  }
   auto in = broadcast_to(astype(vals, dtype, s), shape, s);
   return array(shape, dtype, std::make_unique<Full>(to_stream(s)), {in});
 }
@@ -217,7 +228,7 @@ array ones_like(const array& a, StreamOrDevice s /* = {} */) {
 
 array eye(int n, int m, int k, Dtype dtype, StreamOrDevice s /* = {} */) {
   if (n <= 0 || m <= 0) {
-    throw std::invalid_argument("N and M must be positive integers.");
+    throw std::invalid_argument("[eye] N and M must be positive integers.");
   }
   array result = zeros({n, m}, dtype, s);
   if (k >= m || -k >= n) {
@@ -245,7 +256,7 @@ array tri(int n, int m, int k, Dtype type, StreamOrDevice s /* = {} */) {
   return astype(greater_equal(l, r, s), type, s);
 }
 
-array tril(array x, int k, StreamOrDevice s /* = {} */) {
+array tril(array x, int k /* = 0 */, StreamOrDevice s /* = {} */) {
   if (x.ndim() < 2) {
     throw std::invalid_argument("[tril] array must be at least 2-D");
   }
@@ -253,7 +264,7 @@ array tril(array x, int k, StreamOrDevice s /* = {} */) {
   return where(mask, x, zeros_like(x, s), s);
 }
 
-array triu(array x, int k, StreamOrDevice s /* = {} */) {
+array triu(array x, int k /* = 0 */, StreamOrDevice s /* = {} */) {
   if (x.ndim() < 2) {
     throw std::invalid_argument("[triu] array must be at least 2-D");
   }
@@ -662,26 +673,27 @@ array concatenate(
     int axis,
     StreamOrDevice s /* = {} */) {
   if (arrays.size() == 0) {
-    throw std::invalid_argument("No arrays provided for concatenation");
+    throw std::invalid_argument(
+        "[concatenate] No arrays provided for concatenation");
   }
 
   // Normalize the given axis
   auto ax = axis < 0 ? axis + arrays[0].ndim() : axis;
   if (ax < 0 || ax >= arrays[0].ndim()) {
     std::ostringstream msg;
-    msg << "Invalid axis (" << axis << ") passed to concatenate"
+    msg << "[concatenate] Invalid axis (" << axis << ") passed to concatenate"
         << " for array with shape " << arrays[0].shape() << ".";
     throw std::invalid_argument(msg.str());
   }
 
   auto throw_invalid_shapes = [&]() {
     std::ostringstream msg;
-    msg << "All the input array dimensions must match exactly except"
-        << " for the concatenation axis. However, the provided shapes are ";
+    msg << "[concatenate] All the input array dimensions must match exactly "
+        << "except for the concatenation axis. However, the provided shapes are ";
     for (auto& a : arrays) {
       msg << a.shape() << ", ";
     }
-    msg << "and the concatenation axis is " << axis;
+    msg << "and the concatenation axis is " << axis << ".";
     throw std::invalid_argument(msg.str());
   };
 
@@ -690,6 +702,13 @@ array concatenate(
   // Make the output shape and validate that all arrays have the same shape
   // except for the concatenation axis.
   for (auto& a : arrays) {
+    if (a.ndim() != shape.size()) {
+      std::ostringstream msg;
+      msg << "[concatenate] All the input arrays must have the same number of "
+          << "dimensions. However, got arrays with dimensions " << shape.size()
+          << " and " << a.ndim() << ".";
+      throw std::invalid_argument(msg.str());
+    }
     for (int i = 0; i < a.ndim(); i++) {
       if (i == ax) {
         continue;
@@ -1105,18 +1124,27 @@ array array_equal(
 }
 
 array isnan(const array& a, StreamOrDevice s /* = {} */) {
+  if (is_integral(a.dtype())) {
+    return full(a.shape(), false, bool_, s);
+  }
   return not_equal(a, a, s);
 }
 
 array isinf(const array& a, StreamOrDevice s /* = {} */) {
+  return logical_or(isposinf(a, s), isneginf(a, s), s);
+}
+
+array isposinf(const array& a, StreamOrDevice s /* = {} */) {
+  if (is_integral(a.dtype())) {
+    return full(a.shape(), false, bool_, s);
+  }
   return equal(a, array(std::numeric_limits<float>::infinity(), a.dtype()), s);
 }
 
-array isposinf(const array& a, StreamOrDevice s) {
-  return equal(a, array(std::numeric_limits<float>::infinity(), a.dtype()), s);
-}
-
-array isneginf(const array& a, StreamOrDevice s) {
+array isneginf(const array& a, StreamOrDevice s /* = {} */) {
+  if (is_integral(a.dtype())) {
+    return full(a.shape(), false, bool_, s);
+  }
   return equal(a, array(-std::numeric_limits<float>::infinity(), a.dtype()), s);
 }
 
@@ -1135,11 +1163,43 @@ array allclose(
     const array& b,
     double rtol /* = 1e-5 */,
     double atol /* = 1e-8 */,
+    bool equal_nan /* = false */,
+    StreamOrDevice s /* = {}*/) {
+  return all(isclose(a, b, rtol, atol, equal_nan, s), s);
+}
+
+array isclose(
+    const array& a,
+    const array& b,
+    double rtol /* = 1e-5 */,
+    double atol /* = 1e-8 */,
+    bool equal_nan /* = false */,
     StreamOrDevice s /* = {}*/) {
   // |a - b| <= atol + rtol * |b|
   auto rhs = add(array(atol), multiply(array(rtol), abs(b, s), s), s);
   auto lhs = abs(subtract(a, b, s), s);
-  return all(less_equal(lhs, rhs, s), s);
+  auto out = less_equal(lhs, rhs, s);
+
+  // Correct the result for infinite values.
+  auto any_inf = logical_or(isinf(a, s), isinf(b, s), s);
+  auto both_inf = logical_or(
+      logical_and(isposinf(a, s), isposinf(b, s), s),
+      logical_and(isneginf(a, s), isneginf(b, s), s),
+      s);
+
+  // Convert all elements where either value is infinite to False.
+  out = logical_and(out, logical_not(any_inf, s), s);
+
+  // Convert all the elements where both values are infinite and of the same
+  // sign to True.
+  out = logical_or(out, both_inf, s);
+
+  if (equal_nan) {
+    auto both_nan = logical_and(isnan(a, s), isnan(b, s), s);
+    out = logical_or(out, both_nan, s);
+  }
+
+  return out;
 }
 
 array all(const array& a, bool keepdims, StreamOrDevice s /* = {}*/) {
@@ -1156,13 +1216,16 @@ array all(
   if (axes.empty()) {
     return astype(a, bool_, s);
   }
-  auto [out_shape, sorted_axes] =
-      compute_reduce_shape(axes, a.shape(), keepdims);
-  return array(
+  auto [out_shape, sorted_axes] = compute_reduce_shape(axes, a.shape());
+  auto out = array(
       out_shape,
       bool_,
       std::make_unique<Reduce>(to_stream(s), Reduce::And, sorted_axes),
       {a});
+  if (!keepdims) {
+    out = squeeze(out, sorted_axes, s);
+  }
+  return out;
 }
 
 array all(
@@ -1187,13 +1250,16 @@ array any(
   if (axes.empty()) {
     return astype(a, bool_, s);
   }
-  auto [out_shape, sorted_axes] =
-      compute_reduce_shape(axes, a.shape(), keepdims);
-  return array(
+  auto [out_shape, sorted_axes] = compute_reduce_shape(axes, a.shape());
+  auto out = array(
       out_shape,
       bool_,
       std::make_unique<Reduce>(to_stream(s), Reduce::Or, sorted_axes),
       {a});
+  if (!keepdims) {
+    out = squeeze(out, sorted_axes, s);
+  }
+  return out;
 }
 
 array any(
@@ -1218,14 +1284,17 @@ array sum(
   if (axes.empty()) {
     return a;
   }
-  auto [out_shape, sorted_axes] =
-      compute_reduce_shape(axes, a.shape(), keepdims);
+  auto [out_shape, sorted_axes] = compute_reduce_shape(axes, a.shape());
   auto out_type = a.dtype() == bool_ ? int32 : a.dtype();
-  return array(
+  auto out = array(
       out_shape,
       out_type,
       std::make_unique<Reduce>(to_stream(s), Reduce::Sum, sorted_axes),
       {a});
+  if (!keepdims) {
+    out = squeeze(out, sorted_axes, s);
+  }
+  return out;
 }
 
 array sum(
@@ -1276,11 +1345,18 @@ array var(
     bool keepdims /* = false */,
     int ddof /* = 0*/,
     StreamOrDevice s /* = {}*/) {
-  auto nelements = compute_number_of_elements(a, axes);
   auto dtype = at_least_float(a.dtype());
-  auto mu = mean(a, axes, true, s);
-  auto S = sum(square(subtract(a, mu, s), s), axes, keepdims, s);
-  return multiply(S, array(1.0 / (nelements - ddof), dtype), s);
+  auto mu2 = square(mean(a, axes, keepdims, s), s);
+  auto a2 = mean(square(a, s), axes, keepdims, s);
+  auto v = subtract(a2, mu2, s);
+
+  if (ddof != 0) {
+    auto nelements = compute_number_of_elements(a, axes);
+    float factor = nelements / (nelements - ddof);
+    v = multiply(v, array(factor, dtype), s);
+  }
+
+  return v;
 }
 
 array var(
@@ -1306,13 +1382,16 @@ array prod(
   if (axes.empty()) {
     return a;
   }
-  auto [out_shape, sorted_axes] =
-      compute_reduce_shape(axes, a.shape(), keepdims);
-  return array(
+  auto [out_shape, sorted_axes] = compute_reduce_shape(axes, a.shape());
+  auto out = array(
       out_shape,
       a.dtype(),
       std::make_unique<Reduce>(to_stream(s), Reduce::Prod, sorted_axes),
       {a});
+  if (!keepdims) {
+    out = squeeze(out, sorted_axes, s);
+  }
+  return out;
 }
 
 array prod(
@@ -1340,13 +1419,16 @@ array max(
   if (axes.empty()) {
     return a;
   }
-  auto [out_shape, sorted_axes] =
-      compute_reduce_shape(axes, a.shape(), keepdims);
-  return array(
+  auto [out_shape, sorted_axes] = compute_reduce_shape(axes, a.shape());
+  auto out = array(
       out_shape,
       a.dtype(),
       std::make_unique<Reduce>(to_stream(s), Reduce::Max, sorted_axes),
       {a});
+  if (!keepdims) {
+    out = squeeze(out, sorted_axes, s);
+  }
+  return out;
 }
 
 array max(
@@ -1374,13 +1456,16 @@ array min(
   if (axes.empty()) {
     return a;
   }
-  auto [out_shape, sorted_axes] =
-      compute_reduce_shape(axes, a.shape(), keepdims);
-  return array(
+  auto [out_shape, sorted_axes] = compute_reduce_shape(axes, a.shape());
+  auto out = array(
       out_shape,
       a.dtype(),
       std::make_unique<Reduce>(to_stream(s), Reduce::Min, sorted_axes),
       {a});
+  if (!keepdims) {
+    out = squeeze(out, sorted_axes, s);
+  }
+  return out;
 }
 
 array min(
@@ -1409,14 +1494,17 @@ array argmin(
     throw std::invalid_argument(
         "[argmin] Cannot argmin reduce zero size array.");
   }
-  auto [out_shape, sorted_axes] =
-      compute_reduce_shape({axis}, a.shape(), keepdims);
-  return array(
+  auto [out_shape, sorted_axes] = compute_reduce_shape({axis}, a.shape());
+  auto out = array(
       out_shape,
       uint32,
       std::make_unique<ArgReduce>(
           to_stream(s), ArgReduce::ArgMin, sorted_axes[0]),
       {a});
+  if (!keepdims) {
+    out = squeeze(out, sorted_axes, s);
+  }
+  return out;
 }
 
 array argmax(const array& a, bool keepdims, StreamOrDevice s /* = {} */) {
@@ -1437,14 +1525,17 @@ array argmax(
     throw std::invalid_argument(
         "[argmax] Cannot argmax reduce zero size array.");
   }
-  auto [out_shape, sorted_axes] =
-      compute_reduce_shape({axis}, a.shape(), keepdims);
-  return array(
+  auto [out_shape, sorted_axes] = compute_reduce_shape({axis}, a.shape());
+  auto out = array(
       out_shape,
       uint32,
       std::make_unique<ArgReduce>(
           to_stream(s), ArgReduce::ArgMax, sorted_axes[0]),
       {a});
+  if (!keepdims) {
+    out = squeeze(out, sorted_axes, s);
+  }
+  return out;
 }
 
 /** Returns a sorted copy of the flattened array. */
@@ -2600,9 +2691,40 @@ inline std::vector<int> conv_out_shape(
   std::vector<int> out_shape(in_shape.size());
   int i = 0;
   out_shape[i++] = N;
+
   for (; i < in_shape.size() - 1; i++) {
+    if (pads[i - 1] < 0) {
+      std::ostringstream msg;
+      msg << "[conv] Padding sizes must be non-negative."
+          << " Got padding " << pads << ".";
+      throw std::invalid_argument(msg.str());
+    }
+
+    if (strides[i - 1] <= 0) {
+      std::ostringstream msg;
+      msg << "[conv] Stride sizes must be positive."
+          << " Got strides " << strides << ".";
+      throw std::invalid_argument(msg.str());
+    }
+
+    if (dilation[i - 1] <= 0) {
+      std::ostringstream msg;
+      msg << "[conv] Dilation sizes must be positive."
+          << " Got dilation " << dilation << ".";
+      throw std::invalid_argument(msg.str());
+    }
+
     out_shape[i] = conv_out_axis_size(
         in_shape[i], wt_shape[i], strides[i - 1], pads[i - 1], dilation[i - 1]);
+
+    if (out_shape[i] <= 0) {
+      std::ostringstream msg;
+      msg << "[conv] Spatial dimensions of input after padding "
+          << " cannot be smaller than weight spatial dimensions."
+          << " Got input with shape " << in_shape << " and padding " << pads
+          << " for weight of shape " << wt_shape << ".";
+      throw std::invalid_argument(msg.str());
+    }
   }
   out_shape[i] = O;
 
@@ -2833,7 +2955,7 @@ std::tuple<array, array, array> quantize(
     int group_size /* = 64 */,
     int bits /* = 4 */,
     StreamOrDevice s /* = {} */) {
-  if (group_size != 64 && group_size != 128) {
+  if (group_size != 32 && group_size != 64 && group_size != 128) {
     std::ostringstream msg;
     msg << "[quantize] The requested group size " << group_size
         << " is not supported. The supported group sizes are 64 and 128.";
@@ -2906,6 +3028,16 @@ array dequantize(
     int group_size /* = 64 */,
     int bits /* = 4 */,
     StreamOrDevice s /* = {} */) {
+  if (bits <= 0) {
+    std::ostringstream msg;
+    msg << "[dequantize] Invalid value for bits: " << bits;
+    throw std::invalid_argument(msg.str());
+  }
+  if (group_size <= 0) {
+    std::ostringstream msg;
+    msg << "[dequantize] Invalid value for group_size: " << group_size;
+    throw std::invalid_argument(msg.str());
+  }
   if (w.ndim() != 2 || scales.ndim() != 2 || biases.ndim() != 2) {
     throw std::invalid_argument("[dequantize] Only matrices supported for now");
   }
@@ -3151,4 +3283,102 @@ array addmm(
   return out;
 }
 
+array diagonal(
+    const array& a,
+    int offset /* = 0 */,
+    int axis1 /* = 0 */,
+    int axis2 /* = 1 */,
+    StreamOrDevice s /* = {} */
+) {
+  int ndim = a.ndim();
+  if (ndim < 2) {
+    std::ostringstream msg;
+    msg << "[diagonal] Array must have at least two dimensions, but got "
+        << ndim << " dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+
+  auto ax1 = (axis1 < 0) ? axis1 + ndim : axis1;
+  if (ax1 < 0 || ax1 >= ndim) {
+    std::ostringstream msg;
+    msg << "[diagonal] Invalid axis1 " << axis1 << " for array with " << ndim
+        << " dimensions.";
+    throw std::out_of_range(msg.str());
+  }
+
+  auto ax2 = (axis2 < 0) ? axis2 + ndim : axis2;
+  if (ax2 < 0 || ax2 >= ndim) {
+    std::ostringstream msg;
+    msg << "[diagonal] Invalid axis2 " << axis2 << " for array with " << ndim
+        << " dimensions.";
+    throw std::out_of_range(msg.str());
+  }
+
+  if (ax1 == ax2) {
+    throw std::invalid_argument(
+        "[diagonal] axis1 and axis2 cannot be the same axis");
+  }
+
+  auto off1 = std::max(-offset, 0);
+  auto off2 = std::max(offset, 0);
+
+  auto diag_size = std::min(a.shape(ax1) - off1, a.shape(ax2) - off2);
+  diag_size = std::max(diag_size, 0);
+
+  std::vector<array> indices = {
+      arange(off1, off1 + diag_size, s), arange(off2, off2 + diag_size, s)};
+
+  std::vector<int> slice_sizes = a.shape();
+  slice_sizes[ax1] = 1;
+  slice_sizes[ax2] = 1;
+
+  auto out = gather(a, indices, {ax1, ax2}, slice_sizes, s);
+  return moveaxis(squeeze(out, {ax1 + 1, ax2 + 1}, s), 0, -1, s);
+}
+
+array diag(const array& a, int k /* = 0 */, StreamOrDevice s /* = {} */) {
+  if (a.ndim() == 1) {
+    int a_size = a.size();
+    int n = a_size + std::abs(k);
+    auto res = zeros({n, n}, a.dtype(), s);
+
+    std::vector<array> indices;
+    auto s1 = std::max(0, -k);
+    auto s2 = std::max(0, k);
+    indices.push_back(arange(s1, a_size + s1, uint32, s));
+    indices.push_back(arange(s2, a_size + s2, uint32, s));
+
+    return scatter(res, indices, reshape(a, {a_size, 1, 1}, s), {0, 1}, s);
+  } else if (a.ndim() == 2) {
+    return diagonal(a, k, 0, 1, s);
+  } else {
+    std::ostringstream msg;
+    msg << "[diag] array must be 1-D or 2-D, got array with " << a.ndim()
+        << " dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+}
+
+std::vector<array> depends(
+    const std::vector<array>& inputs,
+    const std::vector<array>& dependencies) {
+  std::vector<array> all_inputs = inputs;
+  all_inputs.insert(all_inputs.end(), dependencies.begin(), dependencies.end());
+
+  // Compute the stream. Maybe do it in a smarter way at some point in the
+  // future.
+  Stream s = (inputs[0].has_primitive()) ? inputs[0].primitive().stream()
+                                         : to_stream({});
+  // Make the output info
+  std::vector<std::vector<int>> shapes;
+  std::vector<Dtype> dtypes;
+  for (const auto& in : inputs) {
+    shapes.emplace_back(in.shape());
+    dtypes.emplace_back(in.dtype());
+  }
+
+  return array::make_arrays(
+      shapes, dtypes, std::make_shared<Depends>(to_stream(s)), all_inputs);
+}
+
 } // namespace mlx::core

mlx/ops.h

@@ -1,14 +1,13 @@
-// Copyright © 2023 Apple Inc.
+// Copyright © 2023-2024 Apple Inc.
 
 #pragma once
 
 #include <optional>
 #include <variant>
 
-#include "array.h"
-#include "device.h"
-#include "io/load.h"
-#include "stream.h"
+#include "mlx/array.h"
+#include "mlx/device.h"
+#include "mlx/stream.h"
 
 namespace mlx::core {
 
@@ -124,8 +123,8 @@ inline array tri(int n, Dtype type, StreamOrDevice s = {}) {
   return tri(n, n, 0, type, s);
 }
 
-array tril(array x, int k, StreamOrDevice s = {});
-array triu(array x, int k, StreamOrDevice s = {});
+array tril(array x, int k = 0, StreamOrDevice s = {});
+array triu(array x, int k = 0, StreamOrDevice s = {});
 
 /** array manipulation */
 
@@ -405,6 +404,17 @@ array allclose(
     const array& b,
     double rtol = 1e-5,
     double atol = 1e-8,
+    bool equal_nan = false,
+    StreamOrDevice s = {});
+
+/** Returns a boolean array where two arrays are element-wise equal within the
+ * specified tolerance. */
+array isclose(
+    const array& a,
+    const array& b,
+    double rtol = 1e-5,
+    double atol = 1e-8,
+    bool equal_nan = false,
     StreamOrDevice s = {});
 
 /**
@@ -1040,20 +1050,6 @@ array conv2d(
     int groups = 1,
     StreamOrDevice s = {});
 
-/** Serialization operations */
-
-/** Save array to out stream in .npy format */
-void save(std::shared_ptr<io::Writer> out_stream, array a);
-
-/** Save array to file in .npy format */
-void save(const std::string& file, array a);
-
-/** Load array from reader in .npy format */
-array load(std::shared_ptr<io::Reader> in_stream, StreamOrDevice s = {});
-
-/** Load array from file in .npy format */
-array load(const std::string& file, StreamOrDevice s = {});
-
 /** Quantized matmul multiplies x with a quantized matrix w*/
 array quantized_matmul(
     const array& x,
@@ -1100,28 +1096,6 @@ array outer(const array& a, const array& b, StreamOrDevice s = {});
 /** Compute the inner product of two vectors. */
 array inner(const array& a, const array& b, StreamOrDevice s = {});
 
-/** Load array map from .safetensors file format */
-std::unordered_map<std::string, array> load_safetensors(
-    std::shared_ptr<io::Reader> in_stream,
-    StreamOrDevice s = {});
-std::unordered_map<std::string, array> load_safetensors(
-    const std::string& file,
-    StreamOrDevice s = {});
-
-void save_safetensors(
-    std::shared_ptr<io::Writer> in_stream,
-    std::unordered_map<std::string, array>);
-void save_safetensors(
-    const std::string& file,
-    std::unordered_map<std::string, array>);
-
-/** Load array map from .gguf file format */
-std::unordered_map<std::string, array> load_gguf(
-    const std::string& file,
-    StreamOrDevice s = {});
-
-void save_gguf(std::string file, std::unordered_map<std::string, array> a);
-
 /** Compute D = beta * C + alpha * (A @ B) */
 array addmm(
     array c,
@@ -1130,4 +1104,25 @@ array addmm(
     const float& alpha = 1.f,
     const float& beta = 1.f,
     StreamOrDevice s = {});
+
+/** Extract a diagonal or construct a diagonal array */
+array diagonal(
+    const array& a,
+    int offset = 0,
+    int axis1 = 0,
+    int axis2 = 1,
+    StreamOrDevice s = {});
+
+/** Extract diagonal from a 2d array or create a diagonal matrix. */
+array diag(const array& a, int k = 0, StreamOrDevice s = {});
+
+/**
+ * Implements the identity function but allows injecting dependencies to other
+ * arrays. This ensures that these other arrays will have been computed
+ * when the outputs of this function are computed.
+ */
+std::vector<array> depends(
+    const std::vector<array>& inputs,
+    const std::vector<array>& dependencies);
+
 } // namespace mlx::core

mlx/primitives.cpp

@@ -1,4 +1,4 @@
-// Copyright © 2023 Apple Inc.
+// Copyright © 2023-2024 Apple Inc.
 #include <algorithm>
 #include <cassert>
 #include <cmath>
@@ -360,6 +360,20 @@ bool ArgReduce::is_equivalent(const Primitive& other) const {
   return reduce_type_ == r_other.reduce_type_ && axis_ == r_other.axis_;
 }
 
+std::pair<std::vector<array>, std::vector<int>> ArgReduce::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  int reduce_ax = axis_ + (axis_ >= axes[0]);
+  auto& in = inputs[0];
+  std::vector<array> out;
+  if (reduce_type_ == ArgReduce::ArgMin) {
+    out.push_back(argmin(in, reduce_ax, true, stream()));
+  } else {
+    out.push_back(argmax(in, reduce_ax, true, stream()));
+  }
+  return {out, axes};
+}
+
 std::pair<std::vector<array>, std::vector<int>> ArgSort::vmap(
     const std::vector<array>& inputs,
     const std::vector<int>& axes) {
@@ -797,6 +811,43 @@ std::pair<std::vector<array>, std::vector<int>> Cosh::vmap(
   return {{cosh(inputs[0], stream())}, axes};
 }
 
+std::vector<array> CustomVJP::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>& outputs) {
+  std::vector<array> inputs(primals.begin(), primals.end() - outputs.size());
+  auto all_vjps = vjp_fun_(inputs, cotangents, outputs);
+  for (const auto& cot : cotangents) {
+    all_vjps.emplace_back(cot);
+  }
+
+  std::vector<array> vjps;
+  vjps.reserve(argnums.size());
+  for (auto arg : argnums) {
+    vjps.push_back(all_vjps[arg]);
+  }
+
+  return vjps;
+}
+
+std::vector<array> Depends::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>& outputs) {
+  std::vector<array> vjps;
+
+  for (auto arg : argnums) {
+    if (arg < cotangents.size()) {
+      vjps.push_back(cotangents[arg]);
+    } else {
+      vjps.push_back(zeros_like(primals[arg]));
+    }
+  }
+  return vjps;
+}
+
 std::vector<array> Divide::vjp(
     const std::vector<array>& primals,
     const std::vector<array>& cotangents,
@@ -1852,7 +1903,12 @@ std::vector<array> Power::vjp(
   for (auto arg : argnums) {
     if (arg == 0) {
       vjps.push_back(multiply(
-          outputs[0], divide(primals[1], primals[0], stream()), stream()));
+          power(
+              primals[0],
+              subtract(primals[1], array(1, primals[0].dtype()), stream()),
+              stream()),
+          primals[1],
+          stream()));
     } else {
       vjps.push_back(multiply(log(primals[0], stream()), outputs[0], stream()));
     }
@@ -2110,7 +2166,36 @@ std::vector<array> Reduce::vjp(
 std::pair<std::vector<array>, std::vector<int>> Reduce::vmap(
     const std::vector<array>& inputs,
     const std::vector<int>& axes) {
-  throw std::runtime_error("Reduce::vmap not yet implemented.");
+  auto ax = axes[0];
+  auto reduce_axes = axes_;
+  for (auto& rax : reduce_axes) {
+    if (rax >= ax) {
+      rax++;
+    }
+  }
+  auto& in = inputs[0];
+  std::vector<array> out;
+  switch (reduce_type_) {
+    case Reduce::And:
+      out.push_back(all(in, reduce_axes, true, stream()));
+      break;
+    case Reduce::Or:
+      out.push_back(any(in, reduce_axes, true, stream()));
+      break;
+    case Reduce::Sum:
+      out.push_back(sum(in, reduce_axes, true, stream()));
+      break;
+    case Reduce::Prod:
+      out.push_back(prod(in, reduce_axes, true, stream()));
+      break;
+    case Reduce::Min:
+      out.push_back(min(in, reduce_axes, true, stream()));
+      break;
+    case Reduce::Max:
+      out.push_back(max(in, reduce_axes, true, stream()));
+      break;
+  }
+  return {out, axes};
 }
 
 bool Reduce::is_equivalent(const Primitive& other) const {

mlx/primitives.h

@@ -1,4 +1,4 @@
-// Copyright © 2023 Apple Inc.
+// Copyright © 2023-2024 Apple Inc.
 
 #pragma once
 
@@ -341,6 +341,7 @@ class ArgReduce : public UnaryPrimitive {
   void eval_cpu(const std::vector<array>& inputs, array& out) override;
   void eval_gpu(const std::vector<array>& inputs, array& out) override;
 
+  DEFINE_VMAP()
   DEFINE_PRINT(ArgReduce)
   bool is_equivalent(const Primitive& other) const override;
 
@@ -552,6 +553,60 @@ class Cosh : public UnaryPrimitive {
   void eval(const std::vector<array>& inputs, array& out);
 };
 
+class CustomVJP : public Primitive {
+ public:
+  explicit CustomVJP(
+      Stream stream,
+      std::function<std::vector<array>(
+          const std::vector<array>&,
+          const std::vector<array>&,
+          const std::vector<array>&)> fun)
+      : Primitive(stream), vjp_fun_(std::move(fun)) {}
+
+  void eval_cpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+  void eval_gpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+
+  std::vector<array> vjp(
+      const std::vector<array>& primals,
+      const std::vector<array>& cotan,
+      const std::vector<int>& argnums,
+      const std::vector<array>& outputs) override;
+
+  DEFINE_PRINT(CustomVJP);
+
+ private:
+  void eval(const std::vector<array>& inputs, std::vector<array>& outputs);
+
+  std::function<std::vector<array>(
+      const std::vector<array>&,
+      const std::vector<array>&,
+      const std::vector<array>&)>
+      vjp_fun_;
+};
+
+class Depends : public Primitive {
+ public:
+  explicit Depends(Stream stream) : Primitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+  void eval_gpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+
+  std::vector<array> vjp(
+      const std::vector<array>& primals,
+      const std::vector<array>& cotan,
+      const std::vector<int>& argnums,
+      const std::vector<array>& outputs) override;
+
+  DEFINE_PRINT(Depends);
+
+ private:
+  void eval(const std::vector<array>& inputs, std::vector<array>& outputs);
+};
+
 class Divide : public UnaryPrimitive {
  public:
   explicit Divide(Stream stream) : UnaryPrimitive(stream){};
@@ -1602,4 +1657,20 @@ class Transpose : public UnaryPrimitive {
   void eval(const std::vector<array>& inputs, array& out);
 };
 
+/* QR Factorization primitive. */
+class QRF : public Primitive {
+ public:
+  explicit QRF(Stream stream) : Primitive(stream){};
+
+  void eval_cpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+  void eval_gpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+
+  DEFINE_PRINT(QRF)
+
+ private:
+  void eval(const std::vector<array>& inputs, std::vector<array>& outputs);
+};
+
 } // namespace mlx::core

mlx/transforms.cpp

@@ -1,7 +1,6 @@
-// Copyright © 2023 Apple Inc.
+// Copyright © 2023-2024 Apple Inc.
 #include <algorithm>
 #include <future>
-#include <map>
 #include <numeric>
 #include <set>
 #include <sstream>
@@ -35,171 +34,8 @@ class Synchronizer : public Primitive {
 // are currently under a function transformation.
 int detail::InTracing::tracing_counter{0};
 
-void simplify(const std::vector<array>& outputs) {
-  // Some notes about how this function works
-  //
-  // Step 1: Traverse the graph and build a tape. During the graph
-  // traversal we:
-  //      - Build a map of inputs to their parents.
-  //      - Record scalar inputs in a map in order to fuse them.
-  // Step 2: Process the tape. A node in the tape has inputs and outputs.
-  //      - Scalar inputs are replaced with their canonical scalar
-  //      - We check each inputs output nodes. Every output node that matches
-  //        the current node gets fused into the current node.
-  std::function<void(const array&)> recurse;
-  std::queue<array> tape;
-  std::unordered_set<std::uintptr_t> cache;
-  std::unordered_map<std::uintptr_t, std::vector<std::pair<array, int>>>
-      parents_map;
-
-  // Helpers to identify identical scalars
-  std::map<std::pair<uint64_t, Dtype::Val>, array> scalars;
-  auto is_scalar = [](const array& a) {
-    return a.is_evaled() && a.ndim() == 0;
-  };
-  auto get_scalar_rep = [](const array& a) {
-    uint64_t v = 0;
-    int dtype;
-    switch (a.dtype().size) {
-      case 1:
-        v = *a.data<uint8_t>();
-        break;
-      case 4:
-        v = *a.data<uint32_t>();
-        break;
-      case 8:
-        v = *a.data<uint64_t>();
-        break;
-    }
-    return std::make_pair(v, a.dtype().val);
-  };
-
-  // DFS the graph to build the tape, and log parents and scalars
-  recurse = [&](const array& a) {
-    auto id = a.id();
-    if (cache.find(id) != cache.end()) {
-      return;
-    }
-    for (int i = 0; i < a.inputs().size(); i++) {
-      auto& in = a.inputs()[i];
-      parents_map[in.id()].push_back({a, i});
-      for (auto& s : a.siblings()) {
-        parents_map[in.id()].push_back({s, i});
-      }
-      recurse(in);
-    }
-    cache.insert(id);
-    for (auto& s : a.siblings()) {
-      cache.insert(s.id());
-    }
-
-    tape.push(a);
-    if (is_scalar(a)) {
-      scalars.insert({get_scalar_rep(a), a});
-    }
-  };
-  for (auto& a : outputs) {
-    recurse(a);
-  }
-
-  // Helper that fuses two arrays in the graph by setting the parents of the
-  // source to point to the destination
-  auto fuse = [&](array& dst, array& src) {
-    // Canonicalize the order of the primitives outputs
-    auto sources = src.outputs();
-    auto dests = dst.outputs();
-    // For each src parent, point it to the corresponding dest
-    for (int i = 0; i < sources.size(); ++i) {
-      auto src_parents = parents_map.find(sources[i].id());
-      if (src_parents == parents_map.end()) {
-        continue;
-      }
-      auto& pairs = parents_map[dests[i].id()];
-      for (auto& parent : src_parents->second) {
-        parent.first.inputs()[parent.second] = dests[i];
-        pairs.push_back(parent);
-      }
-      // Remove the source from the map to avoid fusing with it again
-      parents_map.erase(src_parents);
-    }
-  };
-
-  // Depth-1 array equivalence check.
-  auto array_equivalent = [](const array& a, const array& b) {
-    if (!a.has_primitive() || !b.has_primitive()) {
-      return false;
-    }
-    if (a.primitive_id() == b.primitive_id()) {
-      return false;
-    }
-    const auto& pa = a.primitive();
-    const auto& pb = b.primitive();
-    if (typeid(pa) != typeid(pb)) {
-      return false;
-    }
-
-    if (a.inputs().size() != b.inputs().size()) {
-      return false;
-    }
-
-    for (int i = 0; i < a.inputs().size(); i++) {
-      if (a.inputs()[i].id() != b.inputs()[i].id()) {
-        return false;
-      }
-    }
-
-    return pa.is_equivalent(pb);
-  };
-
-  // Walk the graph
-  while (!tape.empty()) {
-    auto arr = std::move(tape.front());
-    tape.pop();
-
-    // Check if we can fuse scalars
-    if (is_scalar(arr)) {
-      auto scalar = scalars.find(get_scalar_rep(arr));
-      if (scalar->second.id() != arr.id()) {
-        fuse(scalar->second, arr);
-        arr = scalar->second;
-      }
-    }
-
-    // Helper to check if we can fuse the parents of the
-    // given array
-    auto maybe_fuse_parents = [&](auto& a) {
-      auto parents = parents_map.find(a.id());
-      if (parents != parents_map.end()) {
-        auto N = parents->second.size();
-        std::vector<bool> mask(N, false);
-        for (int i = 0; i < N; i++) {
-          if (mask[i]) {
-            continue;
-          }
-          for (int j = i + 1; j < N; j++) {
-            if (mask[j]) {
-              continue;
-            }
-            auto& src = parents->second[j].first;
-            auto& dst = parents->second[i].first;
-            if (src.id() != dst.id() && array_equivalent(src, dst)) {
-              fuse(dst, src);
-              mask[j] = true;
-            }
-          }
-        }
-      }
-    };
-
-    maybe_fuse_parents(arr);
-    for (auto& s : arr.siblings()) {
-      maybe_fuse_parents(s);
-    }
-  }
-}
-
 void eval(const std::vector<array>& outputs) {
-  std::function<void(const array&)> recurse;
+  std::function<void(const array&, bool)> recurse;
   std::queue<array> tape;
   std::unordered_set<std::uintptr_t> cache;
   std::unordered_map<std::uintptr_t, std::shared_future<void>> deps;
@@ -216,21 +52,57 @@ void eval(const std::vector<array>& outputs) {
   auto synchronizer =
       array({}, bool_, std::make_unique<Synchronizer>(stream), outputs);
 
-  recurse = [&](const array& a) {
+  recurse = [&](const array& a, bool largest_branch_first) {
     auto id = a.id();
     if (cache.find(id) != cache.end()) {
       return;
     }
-    for (auto in : a.inputs()) {
-      recurse(in);
-      // If one of the inputs is being computed on a different
-      // stream, we need to manage the dependency.
+
+    // If the input is being computed on a different stream, we need to manage
+    // the dependency.
+    auto check_dependency = [&](const array& in) {
       if (!in.is_evaled()) {
         if (a.primitive().stream() != in.primitive().stream()) {
           deps.insert({in.primitive_id(), std::shared_future<void>{}});
         }
       }
+    };
+
+    // Recurse to the largest or smallest branch first.
+    size_t num_inputs = a.inputs().size();
+    if (num_inputs == 1) {
+      auto& in = a.inputs()[0];
+      recurse(in, true);
+      check_dependency(in);
+    } else if (num_inputs == 2) {
+      auto depth_1 = a.inputs()[0].graph_depth();
+      auto depth_2 = a.inputs()[1].graph_depth();
+      auto& in1 = a.inputs()[static_cast<int>(
+          !((depth_1 > depth_2) == largest_branch_first))];
+      auto& in2 = a.inputs()[static_cast<int>(
+          ((depth_1 > depth_2) == largest_branch_first))];
+      recurse(in1, true);
+      check_dependency(in1);
+      recurse(in2, true);
+      check_dependency(in2);
+    } else if (num_inputs > 2) {
+      std::vector<int> recursion_order(a.inputs().size());
+      std::iota(recursion_order.begin(), recursion_order.end(), 0);
+      std::sort(
+          recursion_order.begin(),
+          recursion_order.end(),
+          [&a, largest_branch_first](int i, int j) {
+            auto depth_i = a.inputs()[i].graph_depth();
+            auto depth_j = a.inputs()[j].graph_depth();
+            return largest_branch_first ? depth_i > depth_j : depth_j < depth_i;
+          });
+      for (int idx : recursion_order) {
+        auto& in = a.inputs()[idx];
+        recurse(in, true);
+        check_dependency(in);
+      }
     }
+
     cache.insert(id);
     for (auto& s : a.siblings()) {
       cache.insert(s.id());
@@ -244,7 +116,7 @@ void eval(const std::vector<array>& outputs) {
     }
   };
 
-  recurse(synchronizer);
+  recurse(synchronizer, false);
   uintptr_t synch_id = synchronizer.primitive_id();
   deps.insert({synch_id, std::shared_future<void>{}});
 
@@ -262,6 +134,7 @@ void eval(const std::vector<array>& outputs) {
     auto stream = arr.primitive().stream();
     std::vector<std::shared_future<void>> arr_deps;
     for (auto& in : arr.inputs()) {
+      // TODO that's a bug
       if (auto it = deps.find(in.primitive_id()); it != deps.end()) {
         arr_deps.push_back(it->second);
       }
@@ -337,12 +210,21 @@ std::pair<std::vector<array>, std::vector<array>> vjp(
       }
     }
     if (cotan_index >= cotans.size()) {
-      throw std::invalid_argument(
-          "[vjp] Number of outputs with gradient does not match number of cotangents.");
+      std::ostringstream msg;
+      msg << "[vjp] Number of outputs to compute gradients for ("
+          << outputs.size() << ") does not match number of cotangents ("
+          << cotans.size() << ").";
+      throw std::invalid_argument(msg.str());
     }
     if (out.shape() != cotans[cotan_index].shape()) {
-      throw std::invalid_argument(
-          "[vjp] Output shape does not match shape of cotangent.");
+      std::ostringstream msg;
+      msg << "[vjp] Output shape " << out.shape()
+          << " does not match cotangent shape " << cotans[cotan_index].shape()
+          << ".";
+      if (outputs.size() == 1 && out.size() == 1) {
+        msg << " If you are using grad your function must return a scalar.";
+      }
+      throw std::invalid_argument(msg.str());
     }
     output_cotan_pairs.emplace_back(i, cotan_index++);
   }
@@ -666,9 +548,8 @@ std::pair<std::vector<array>, std::vector<array>> vmap_trace(
         "[vmap] The number of in axes must match the number of inputs.");
   }
 
-  // Run the function on placeholder inputs
-  // to get the original graph
-  std::vector<array> s_inputs;
+  // Some error checking and get the vmap axis size
+  size_t vmap_ax_size;
   for (int i = 0; i < inputs.size(); ++i) {
     if (in_axes[i] != -1) {
       if (inputs[i].ndim() == 0) {
@@ -681,7 +562,26 @@ std::pair<std::vector<array>, std::vector<array>> vmap_trace(
             << inputs[i].ndim() << " dimensions.";
         throw std::invalid_argument(msg.str());
       }
+      vmap_ax_size = inputs[i].shape(in_axes[i]);
+    }
+  }
+  // Check that all vmapped axes have the same size
+  for (int i = 0; i < inputs.size(); ++i) {
+    if (in_axes[i] != -1) {
+      if (size_t in_ax = inputs[i].shape(in_axes[i]); vmap_ax_size != in_ax) {
+        std::ostringstream msg;
+        msg << "[vmap] Inconsistent axis sizes: " << in_ax << " and "
+            << vmap_ax_size << ".";
+        throw std::invalid_argument(msg.str());
+      }
+    }
+  }
 
+  // Run the function on placeholder inputs
+  // to get the original graph
+  std::vector<array> s_inputs;
+  for (int i = 0; i < inputs.size(); ++i) {
+    if (in_axes[i] != -1) {
       std::vector<int> shape = inputs[i].shape();
       shape.erase(shape.begin() + in_axes[i]);
       array in(shape, inputs[i].dtype(), nullptr, {});
@@ -867,4 +767,58 @@ std::function<array(const array&)> vmap(
   return [vfun](const array& a) { return vfun({a})[0]; };
 }
 
+std::function<std::vector<array>(const std::vector<array>&)> custom_vjp(
+    std::function<std::vector<array>(const std::vector<array>&)> fun,
+    std::function<std::vector<array>(
+        const std::vector<array>&,
+        const std::vector<array>&,
+        const std::vector<array>&)> fun_vjp) {
+  return [fun = std::move(fun),
+          fun_vjp = std::move(fun_vjp)](const std::vector<array>& args) {
+    // Compute the outputs
+    auto outputs = fun(args);
+    for (auto& out : outputs) {
+      out = stop_gradient(out);
+    }
+
+    // Prepare the inputs to the primitive
+    // We also add the outputs to the primitive so that it can "run" the forward
+    // pass.
+    std::vector<array> inputs = args;
+    inputs.insert(inputs.end(), outputs.begin(), outputs.end());
+
+    // Compute the stream. Maybe do it in a smarter way at some point in the
+    // future.
+    Stream s = (outputs[0].has_primitive()) ? outputs[0].primitive().stream()
+                                            : default_stream(default_device());
+
+    // Make the output info
+    std::vector<std::vector<int>> shapes;
+    std::vector<Dtype> dtypes;
+    for (const auto& out : outputs) {
+      shapes.emplace_back(out.shape());
+      dtypes.emplace_back(out.dtype());
+    }
+
+    return array::make_arrays(
+        shapes,
+        dtypes,
+        std::make_shared<CustomVJP>(to_stream(s), fun_vjp),
+        inputs);
+  };
+}
+
+std::function<std::vector<array>(const std::vector<array>&)> checkpoint(
+    std::function<std::vector<array>(const std::vector<array>&)> fun) {
+  auto vjp_fun = [fun](
+                     const std::vector<array>& primals,
+                     const std::vector<array>& cotangents,
+                     const std::vector<array>& outputs) -> std::vector<array> {
+    auto [__, vjps] = vjp(fun, depends(primals, outputs), cotangents);
+    return vjps;
+  };
+
+  return custom_vjp(fun, vjp_fun);
+}
+
 } // namespace mlx::core

mlx/transforms.h

@@ -1,18 +1,25 @@
-// Copyright © 2023 Apple Inc.
+// Copyright © 2023-2024 Apple Inc.
 
 #pragma once
 
-#include "array.h"
+#include "mlx/array.h"
 
 namespace mlx::core {
 
-/** Fuse equivalent arrays to avoid duplicate execution. */
-void simplify(const std::vector<array>& outputs);
+// Compile takes a function and returns a new function
+std::function<std::vector<array>(const std::vector<array>&)> compile(
+    const std::function<std::vector<array>(const std::vector<array>&)>& fun);
 
-template <typename... Arrays>
-void simplify(Arrays... outputs) {
-  simplify(std::vector<array>{std::forward<Arrays>(outputs)...});
-}
+/** Globally disable compilation.
+ * Setting the environment variable ``MLX_DISABLE_COMPILE`` can also
+ * be used to disable compilation.
+ */
+void disable_compile();
+
+/** Globally enable compilation.
+ * This will override the environment variable ``MLX_DISABLE_COMPILE``.
+ */
+void enable_compile();
 
 void eval(const std::vector<array>& outputs);
 
@@ -184,4 +191,22 @@ std::function<std::vector<array>(const std::vector<array>&)> vmap(
     const std::vector<int>& in_axes = {},
     const std::vector<int>& out_axes = {});
 
+/**
+ * Return the results of calling fun with args but if their vjp is computed it
+ * will be computed by fun_vjp.
+ */
+std::function<std::vector<array>(const std::vector<array>&)> custom_vjp(
+    std::function<std::vector<array>(const std::vector<array>&)> fun,
+    std::function<std::vector<array>(
+        const std::vector<array>&,
+        const std::vector<array>&,
+        const std::vector<array>&)> fun_vjp);
+
+/**
+ * Checkpoint the gradient of a function. Namely, discard all intermediate
+ * state and recalculate it when we need to compute the gradient.
+ */
+std::function<std::vector<array>(const std::vector<array>&)> checkpoint(
+    std::function<std::vector<array>(const std::vector<array>&)> fun);
+
 } // namespace mlx::core

mlx/transforms_impl.h

@@ -1,4 +1,4 @@
-// Copyright © 2023 Apple Inc.
+// Copyright © 2023-2024 Apple Inc.
 
 namespace mlx::core::detail {
 
@@ -14,6 +14,15 @@ std::vector<array> vmap_replace(
     const std::vector<int>& in_axes,
     const std::vector<int>& out_axes);
 
+// This is not part of the general C++ API as calling with a bad id is a bad
+// idea.
+std::function<std::vector<array>(const std::vector<array>&)> compile(
+    const std::function<std::vector<array>(const std::vector<array>&)>& fun,
+    size_t fun_id);
+
+// Erase cached compile functions
+void compile_erase(size_t fun_id);
+
 // Create an InTracing object during tracing operations to signify to the rest
 // of the codebase that we are during tracing so evals should not throw away
 // the graph.

python/mlx/nn/__init__.py

@@ -1,5 +1,5 @@
 # Copyright © 2023 Apple Inc.
 
-from mlx.nn import losses
+from mlx.nn import init, losses
 from mlx.nn.layers import *
 from mlx.nn.utils import value_and_grad

python/mlx/nn/init.py

@@ -0,0 +1,350 @@
+# Copyright © 2023-2024 Apple Inc.
+
+import math
+from typing import Callable, Literal
+
+import mlx.core as mx
+
+
+def constant(
+    value: float, dtype: mx.Dtype = mx.float32
+) -> Callable[[mx.array], mx.array]:
+    r"""An initializer that returns an array filled with ``value``.
+
+    Args:
+        value (float): The value to fill the array with.
+        dtype (Dtype, optional): The data type of the array. Default:
+          ``float32``.
+
+    Returns:
+        Callable[[array], array]: An initializer that returns an array with the
+        same shape as the input, filled with ``value``.
+
+    Example:
+
+        >>> init_fn = nn.init.constant(0.5)
+        >>> init_fn(mx.zeros((2, 2)))
+        array([[0.5, 0.5],
+               [0.5, 0.5]], dtype=float32)
+    """
+
+    def initializer(a: mx.array) -> mx.array:
+        return mx.full(a.shape, value, dtype=dtype)
+
+    return initializer
+
+
+def normal(
+    mean: float = 0.0, std: float = 1.0, dtype: mx.Dtype = mx.float32
+) -> Callable[[mx.array], mx.array]:
+    r"""An initializer that returns samples from a normal distribution.
+
+    Args:
+        mean (float, optional): Mean of the normal distribution. Default:
+          ``0.0``.
+        std (float, optional): Standard deviation of the normal distribution.
+          Default: ``1.0``.
+        dtype (Dtype, optional): The data type of the array. Default:
+          ``float32``.
+
+    Returns:
+        Callable[[array], array]: An initializer that returns an array with the
+        same shape as the input, filled with samples from a normal distribution.
+
+    Example:
+
+        >>> init_fn = nn.init.normal()
+        >>> init_fn(mx.zeros((2, 2)))
+        array([[-0.982273, -0.534422],
+               [0.380709, 0.0645099]], dtype=float32)
+    """
+
+    def initializer(a: mx.array) -> mx.array:
+        return std * mx.random.normal(shape=a.shape, dtype=dtype) + mean
+
+    return initializer
+
+
+def uniform(
+    low: float = 0.0, high: float = 1.0, dtype: mx.Dtype = mx.float32
+) -> Callable[[mx.array], mx.array]:
+    r"""An initializer that returns samples from a uniform distribution.
+
+    Args:
+        low (float, optional): The lower bound of the uniform distribution.
+          Default: ``0.0``.
+        high (float, optional): The upper bound of the uniform distribution.
+          Default: ``1.0``
+        dtype (Dtype, optional): The data type of the array. Default: ``float32``.
+
+    Returns:
+        Callable[[array], array]: An initializer that returns an array
+        with the same shape as the input, filled with samples from a uniform
+        distribution
+
+    Example:
+
+        >>> init_fn = nn.init.uniform(low=0, high=1)
+        >>> init_fn(mx.zeros((2, 2)))
+        array([[0.883935, 0.863726],
+               [0.617261, 0.417497]], dtype=float32)
+    """
+
+    def initializer(a: mx.array) -> mx.array:
+        return mx.random.uniform(low, high, a.shape, dtype=dtype)
+
+    return initializer
+
+
+def identity(dtype: mx.Dtype = mx.float32) -> Callable[[mx.array], mx.array]:
+    r"""An initializer that returns an identity matrix.
+
+    Args:
+        dtype (Dtype, optional): The data type of the array. Defaults:
+          ``float32``.
+
+    Returns:
+        Callable[[array], array]: An initializer that returns an identity
+        matrix with the same shape as the input.
+
+    Example:
+
+        >>> init_fn = nn.init.identity()
+        >>> init_fn(mx.zeros((2, 2)))
+        array([[1, 0],
+               [0, 1]], dtype=float32)
+    """
+
+    def initializer(arr: mx.array) -> mx.array:
+        if arr.ndim != 2 or arr.shape[0] != arr.shape[1]:
+            raise ValueError(
+                f"The input array must be a square matrix but got shape {arr.shape}."
+            )
+        return mx.eye(n=arr.shape[0], dtype=dtype)
+
+    return initializer
+
+
+def _calculate_fan_in_fan_out(x):
+    if x.ndim < 2:
+        raise ValueError(
+            "Glorot / He initialization requires at least 2 dimensional input"
+            f" but input with {x.ndim} dimensions."
+        )
+
+    fan_in = x.shape[-1]
+    fan_out = x.shape[0]
+
+    if x.ndim > 2:
+        receptive_field = 1
+        for d in x.shape[1:-1]:
+            receptive_field *= d
+
+        fan_in = fan_in * receptive_field
+        fan_out = fan_out * receptive_field
+
+    return fan_in, fan_out
+
+
+def glorot_normal(
+    dtype: mx.Dtype = mx.float32,
+) -> Callable[[mx.array, float], mx.array]:
+    r"""A Glorot normal initializer.
+
+    This initializer samples from a normal distribution with a standard
+    deviation computed from the number of input (``fan_in``) and output
+    (``fan_out``) units according to:
+
+    .. math::
+        \sigma = \gamma \sqrt{\frac{2.0}{\text{fan_in} + \text{fan_out}}}
+
+    For more details see the original reference: `Understanding the difficulty
+    of training deep feedforward neural networks
+    <https://proceedings.mlr.press/v9/glorot10a.html>`_
+
+    Args:
+        dtype (Dtype, optional): The data type of the array. Default: ``float32``.
+
+    Returns:
+        Callable[[array, float], array]: An initializer that returns an array
+        with the same shape as the input, filled with samples from the Glorot
+        normal distribution.
+
+    Example:
+
+        >>> init_fn = nn.init.glorot_normal()
+        >>> init_fn(mx.zeros((2, 2)))
+        array([[0.191107, 1.61278],
+               [-0.150594, -0.363207]], dtype=float32)
+        >>> init_fn(mx.zeros((2, 2)), gain=4.0)
+        array([[1.89613, -4.53947],
+               [4.48095, 0.995016]], dtype=float32)
+    """
+
+    def initializer(a: mx.array, gain: float = 1.0) -> mx.array:
+        fan_in, fan_out = _calculate_fan_in_fan_out(a)
+        std = gain * math.sqrt(2.0 / (fan_in + fan_out))
+        return mx.random.normal(shape=a.shape, dtype=dtype) * std
+
+    return initializer
+
+
+def glorot_uniform(
+    dtype: mx.Dtype = mx.float32,
+) -> Callable[[mx.array, float], mx.array]:
+    r"""A Glorot uniform initializer.
+
+    This initializer samples from a uniform distribution with a range
+    computed from the number of input (``fan_in``) and output (``fan_out``)
+    units according to:
+
+    .. math::
+        \sigma = \gamma \sqrt{\frac{6.0}{\text{fan_in} + \text{fan_out}}}
+
+    For more details see the original reference: `Understanding the difficulty
+    of training deep feedforward neural networks
+    <https://proceedings.mlr.press/v9/glorot10a.html>`_
+
+    Args:
+        dtype (Dtype, optional): The data type of the array. Default: ``float32``.
+
+    Returns:
+        Callable[[array, float], array]: An initializer that returns an array
+        with the same shape as the input, filled with samples from the Glorot
+        uniform distribution.
+
+    Example:
+
+        >>> init_fn = nn.init.glorot_uniform()
+        >>> init_fn(mx.zeros((2, 2)))
+        array([[0.223404, -0.890597],
+               [-0.379159, -0.776856]], dtype=float32)
+        >>> init_fn(mx.zeros((2, 2)), gain=4.0)
+        array([[-1.90041, 3.02264],
+               [-0.912766, 4.12451]], dtype=float32)
+    """
+
+    def initializer(a: mx.array, gain: float = 1.0) -> mx.array:
+        fan_in, fan_out = _calculate_fan_in_fan_out(a)
+        limit = gain * math.sqrt(6.0 / (fan_in + fan_out))
+        return mx.random.uniform(-limit, limit, a.shape, dtype=dtype)
+
+    return initializer
+
+
+def he_normal(
+    dtype: mx.Dtype = mx.float32,
+) -> Callable[[mx.array, str, float], mx.array]:
+    r"""Build a He normal initializer.
+
+    This initializer samples from a normal distribution with a standard
+    deviation computed from the number of input (``fan_in``) or output
+    (``fan_out``) units according to:
+
+    .. math::
+        \sigma = \gamma \frac{1}{\sqrt{\text{fan}}}
+
+    where :math:`\text{fan}` is either the number of input units when the
+    ``mode`` is ``"fan_in"`` or output units when the ``mode`` is
+    ``"fan_out"``.
+
+    For more details see the original reference: `Delving Deep into Rectifiers:
+    Surpassing Human-Level Performance on ImageNet Classification
+    <https://arxiv.org/abs/1502.01852>`_
+
+    Args:
+        dtype (Dtype, optional): The data type of the array. Defaults to mx.float32.
+
+    Returns:
+        Callable[[array, str, float], array]: An initializer that returns an
+        array with the same shape as the input, filled with samples from the He
+        normal distribution.
+
+    Example:
+
+        >>> init_fn = nn.init.he_normal()
+        >>> init_fn(mx.zeros((2, 2)))  # uses fan_in
+        array([[-1.25211, 0.458835],
+               [-0.177208, -0.0137595]], dtype=float32)
+        >>> init_fn(mx.zeros((2, 2)), mode="fan_out", gain=5)
+        array([[5.6967, 4.02765],
+               [-4.15268, -2.75787]], dtype=float32)
+    """
+
+    def initializer(
+        a: mx.array,
+        mode: Literal["fan_in", "fan_out"] = "fan_in",
+        gain: float = 1.0,
+    ) -> mx.array:
+        fan_in, fan_out = _calculate_fan_in_fan_out(a)
+        if mode == "fan_in":
+            fan = fan_in
+        elif mode == "fan_out":
+            fan = fan_out
+        else:
+            raise ValueError(f"Invalid mode: {mode}. Valid modes are: fan_in, fan_out")
+
+        std = gain / math.sqrt(fan)
+        return mx.random.normal(shape=a.shape, dtype=dtype) * std
+
+    return initializer
+
+
+def he_uniform(
+    dtype: mx.Dtype = mx.float32,
+) -> Callable[[mx.array, str, float], mx.array]:
+    r"""A He uniform (Kaiming uniform) initializer.
+
+    This initializer samples from a uniform distribution with a range
+    computed from the number of input (``fan_in``) or output (``fan_out``)
+    units according to:
+
+    .. math::
+
+        \sigma = \gamma \sqrt{\frac{3.0}{\text{fan}}}
+
+    where :math:`\text{fan}` is either the number of input units when the
+    ``mode`` is ``"fan_in"`` or output units when the ``mode`` is
+    ``"fan_out"``.
+
+    For more details see the original reference: `Delving Deep into Rectifiers:
+    Surpassing Human-Level Performance on ImageNet Classification
+    <https://arxiv.org/abs/1502.01852>`_
+
+
+    Args:
+        dtype (Dtype, optional): The data type of the array. Default: ``float32``.
+
+    Returns:
+        Callable[[array, str, float], array]: An initializer that returns an
+        array with the same shape as the input, filled with samples from  the
+        He uniform distribution.
+
+    Example:
+
+        >>> init_fn = nn.init.he_uniform()
+        >>> init_fn(mx.zeros((2, 2)))  # uses fan_in
+        array([[0.0300242, -0.0184009],
+               [0.793615, 0.666329]], dtype=float32)
+        >>> init_fn(mx.zeros((2, 2)), mode="fan_out", gain=5)
+        array([[-1.64331, -2.16506],
+               [1.08619, 5.79854]], dtype=float32)
+    """
+
+    def initializer(
+        a: mx.array,
+        mode: Literal["fan_in", "fan_out"] = "fan_in",
+        gain: float = 1.0,
+    ) -> mx.array:
+        fan_in, fan_out = _calculate_fan_in_fan_out(a)
+        if mode == "fan_in":
+            fan = fan_in
+        elif mode == "fan_out":
+            fan = fan_out
+        else:
+            raise ValueError(f"Invalid mode: {mode}. Valid modes are: fan_in, fan_out")
+
+        limit = gain * math.sqrt(3.0 / fan)
+        return mx.random.uniform(-limit, limit, a.shape, dtype=dtype)
+
+    return initializer

python/mlx/nn/layers/__init__.py

@@ -18,6 +18,7 @@ from mlx.nn.layers.activations import (
     SiLU,
     Softmax,
     Softplus,
+    Softshrink,
     Softsign,
     Step,
     Tanh,
@@ -39,6 +40,7 @@ from mlx.nn.layers.activations import (
     silu,
     softmax,
     softplus,
+    softshrink,
     softsign,
     step,
     tanh,

python/mlx/nn/layers/activations.py

@@ -89,6 +89,19 @@ def softsign(x):
     return mx.divide(x, 1 + mx.abs(x))
 
 
+def softshrink(x, lambd: float = 0.5):
+    r"""Applies the Softshrink activation function.
+
+    .. math::
+        \text{softshrink}(x) = \begin{cases}
+        x - \lambda & \text{if } x > \lambda \\
+        x + \lambda & \text{if } x < -\lambda \\
+        0 & \text{otherwise}
+        \end{cases}
+    """
+    return mx.where(mx.abs(x) > lambd, x - mx.sign(x) * lambd, 0)
+
+
 def celu(x, alpha=1.0):
     r"""Applies the Continuously Differentiable Exponential Linear Unit.
 
@@ -173,7 +186,7 @@ def glu(x: mx.array, axis: int = -1) -> mx.array:
         textrm{GLU}(x) = a * \sigma(b)
 
     Args:
-        axis (int): The dimension to split along. Default: ``-1``.
+        axis (int): The dimension to split along. Default: ``-1``
     """
     a, b = mx.split(x, indices_or_sections=2, axis=axis)
     return a * mx.sigmoid(b)
@@ -189,7 +202,7 @@ class GLU(Module):
         textrm{GLU}(x) = a * \sigma(b)
 
     Args:
-        axis (int): The dimension to split along. Default: ``-1``.
+        axis (int): The dimension to split along. Default: ``-1``
     """
 
     def __init__(self, axis: int = -1):
@@ -295,7 +308,7 @@ class ReLU(Module):
     r"""Applies the Rectified Linear Unit.
         Simply ``mx.maximum(x, 0)``.
 
-    See :func:`relu`, for the functional equivalent.
+    See :func:`relu` for the functional equivalent.
     """
 
 
@@ -305,7 +318,7 @@ class LeakyReLU(Module):
     Simply ``mx.maximum(negative_slope * x, x)``.
 
     Args:
-        negative_slope: Controls the angle of the negative slope. Default: 1e-2.
+        negative_slope: Controls the angle of the negative slope. Default: ``1e-2``
     """
 
     def __init__(self, negative_slope=1e-2):
@@ -320,10 +333,10 @@ class ELU(Module):
     r"""Applies the Exponential Linear Unit.
         Simply ``mx.where(x > 0, x, alpha * (mx.exp(x) - 1))``.
 
-    See :func:`elu`, for the functional equivalent.
+    See :func:`elu` for the functional equivalent.
 
     Args:
-        alpha: the :math:`\alpha` value for the ELU formulation. Default: 1.0
+        alpha: the :math:`\alpha` value for the ELU formulation. Default: ``1.0``
     """
 
     def __init__(self, alpha=1.0):
@@ -338,7 +351,7 @@ class ELU(Module):
 class ReLU6(Module):
     r"""Applies the Rectified Linear Unit 6.
 
-    See :func:`relu6`, for the functional equivalent.
+    See :func:`relu6` for the functional equivalent.
     """
 
 
@@ -346,7 +359,7 @@ class ReLU6(Module):
 class Softmax(Module):
     r"""Applies the Softmax function.
 
-    See :func:`softmax`, for the functional equivalent.
+    See :func:`softmax` for the functional equivalent.
     """
 
 
@@ -354,7 +367,7 @@ class Softmax(Module):
 class Softplus(Module):
     r"""Applies the Softplus function.
 
-    See :func:`softplus`, for the functional equivalent.
+    See :func:`softplus` for the functional equivalent.
     """
 
 
@@ -362,19 +375,36 @@ class Softplus(Module):
 class Softsign(Module):
     r"""Applies the Softsign function.
 
-    See :func:`softsign`, for the functional equivalent.
+    See :func:`softsign` for the functional equivalent.
     """
 
 
+class Softshrink(Module):
+    r"""Applies the Softshrink function.
+
+    See :func:`softshrink` for the functional equivalent.
+
+    Args:
+        lambd: the :math:`\lambda` value for Softshrink. Default: ``0.5``
+    """
+
+    def __init__(self, lambd=0.5):
+        super().__init__()
+        self.lambd = lambd
+
+    def __call__(self, x):
+        return softshrink(x, self.lambd)
+
+
 class CELU(Module):
     r"""Applies the Continuously Differentiable Exponential Linear Unit.
         Applies :math:`\max(0, x) + \min(0, \alpha * (\exp(x / \alpha) - 1))`
         element wise.
 
-    See :func:`celu`, for the functional equivalent.
+    See :func:`celu` for the functional equivalent.
 
     Args:
-        alpha: the :math:`\alpha` value for the CELU formulation. Default: 1.0
+        alpha: the :math:`\alpha` value for the CELU formulation. Default: ``1.0``
     """
 
     def __init__(self, alpha=1.0):
@@ -389,7 +419,7 @@ class CELU(Module):
 class SiLU(Module):
     r"""Applies the Sigmoid Linear Unit. Also known as Swish.
 
-    See :func:`silu`, for the functional equivalent.
+    See :func:`silu` for the functional equivalent.
     """
 
 
@@ -397,7 +427,7 @@ class SiLU(Module):
 class LogSoftmax(Module):
     r"""Applies the Log Softmax function.
 
-    See :func:`log_softmax`, for the functional equivalent.
+    See :func:`log_softmax` for the functional equivalent.
     """
 
 
@@ -405,7 +435,7 @@ class LogSoftmax(Module):
 class LogSigmoid(Module):
     r"""Applies the Log Sigmoid function.
 
-    See :func:`log_sigmoid`, for the functional equivalent.
+    See :func:`log_sigmoid` for the functional equivalent.
     """
 
 
@@ -414,11 +444,11 @@ class PReLU(Module):
         Applies :math:`\max(0, x) + a * \min(0, x)` element wise, where :math:`a`
         is an array.
 
-    See :func:`prelu`, for the functional equivalent.
+    See :func:`prelu` for the functional equivalent.
 
     Args:
-        num_parameters: number of :math:`a` to learn. Default: 1
-        init: the initial value of :math:`a`. Default: 0.25
+        num_parameters: number of :math:`a` to learn. Default: ``1``
+        init: the initial value of :math:`a`. Default: ``0.25``
     """
 
     def __init__(self, num_parameters=1, init=0.25):
@@ -482,7 +512,7 @@ def tanh(x):
 class Tanh(Module):
     r"""Applies the hyperbolic tangent function.
 
-    See :func:`tanh`, for the functional equivalent.
+    See :func:`tanh` for the functional equivalent.
     """
 
 
@@ -490,7 +520,7 @@ class Tanh(Module):
 class Hardswish(Module):
     r"""Applies the hardswish function, element-wise.
 
-    See :func:`hardswish`, for the functional equivalent.
+    See :func:`hardswish` for the functional equivalent.
     """
 
 
@@ -522,5 +552,5 @@ class Step(Module):
 class SELU(Module):
     r"""Applies the Scaled Exponential Linear Unit.
 
-    See :func:`selu`, for the functional equivalent.
+    See :func:`selu` for the functional equivalent.
     """

python/mlx/nn/layers/base.py

@@ -96,11 +96,11 @@ class Module(dict):
         strict: bool = True,
     ):
         """
-        Update the model's weights from a ``.npz`` or a list.
+        Update the model's weights from a ``.npz``, a ``.safetensors`` file, or a list.
 
         Args:
             file_or_weights (str or list(tuple(str, mx.array))): The path to
-                the weights ``.npz`` file or a list of pairs of parameter names
+                the weights ``.npz`` file (``.npz`` or ``.safetensors``) or a list of pairs of parameter names
                 and arrays.
             strict (bool, optional): If ``True`` then checks that the provided
               weights exactly match the parameters of the model. Otherwise,
@@ -118,6 +118,9 @@ class Module(dict):
                 # Load from file
                 model.load_weights("weights.npz")
 
+                # Load from .safetensors file
+                model.load_weights("weights.safetensors")
+
                 # Load from list
                 weights = [
                     ("weight", mx.random.uniform(shape=(10, 10))),
@@ -166,9 +169,20 @@ class Module(dict):
 
     def save_weights(self, file: str):
         """
-        Save the model's weights to a ``.npz`` file.
+        Save the model's weights to a file. The saving method is determined by the file extension:
+        - ``.npz`` will use :func:`mx.savez`
+        - ``.safetensors`` will use :func:`mx.save_safetensors`
         """
-        mx.savez(file, **dict(tree_flatten(self.parameters())))
+        params_dict = dict(tree_flatten(self.parameters()))
+
+        if file.endswith(".npz"):
+            mx.savez(file, **params_dict)
+        elif file.endswith(".safetensors"):
+            mx.save_safetensors(file, params_dict)
+        else:
+            raise ValueError(
+                "Unsupported file extension. Use '.npz' or '.safetensors'."
+            )
 
     @staticmethod
     def is_module(value):

python/mlx/nn/layers/dropout.py

@@ -81,7 +81,7 @@ class Dropout2d(Module):
         # Dropout is applied on the whole channel
         # 3D input: (1, 1, C)
         # 4D input: (B, 1, 1, C)
-        mask_shape = x.shape
+        mask_shape = list(x.shape)
         mask_shape[-2] = mask_shape[-3] = 1
 
         mask = mx.random.bernoulli(p=self._p_1, shape=mask_shape)

python/mlx/nn/layers/positional_encoding.py

@@ -104,9 +104,11 @@ class RoPE(Module):
         dtype=mx.float32,
     ):
         if (N, D, offset, base, scale, dtype) != cls._cos_sin_theta_key:
-            D = D // 2
+            half_D = D // 2
             positions = mx.arange(offset, N, dtype=dtype) * scale
-            freqs = mx.exp(-mx.arange(0.0, D, dtype=dtype) * (math.log(base) / D))
+            freqs = mx.exp(
+                -mx.arange(0.0, half_D, dtype=dtype) * (math.log(base) / half_D)
+            )
             theta = mx.reshape(positions, (-1, 1)) * mx.reshape(freqs, (1, -1))
             cls._cos_sin_theta_key = (N, D, offset, base, scale, dtype)
             cls._cos_sin_theta_value = (mx.cos(theta), mx.sin(theta))

python/mlx/nn/layers/transformer.py

@@ -9,6 +9,7 @@ from mlx.nn.layers.base import Module
 from mlx.nn.layers.dropout import Dropout
 from mlx.nn.layers.linear import Linear
 from mlx.nn.layers.normalization import LayerNorm
+from mlx.nn.utils import checkpoint
 
 
 class MultiHeadAttention(Module):
@@ -115,7 +116,7 @@ class TransformerEncoderLayer(Module):
         mlp_dims: Optional[int] = None,
         dropout: float = 0.0,
         activation: Callable[[Any], Any] = relu,
-        norm_first: bool = False,
+        norm_first: bool = True,
     ):
         super().__init__()
         mlp_dims = mlp_dims or dims * 4
@@ -166,7 +167,8 @@ class TransformerEncoder(Module):
         mlp_dims: Optional[int] = None,
         dropout: float = 0.0,
         activation=relu,
-        norm_first: bool = False,
+        norm_first: bool = True,
+        checkpoint: bool = False,
     ):
         super().__init__()
         self.layers = [
@@ -176,9 +178,11 @@ class TransformerEncoder(Module):
             for i in range(num_layers)
         ]
         self.ln = LayerNorm(dims)
+        self.checkpoint = checkpoint
 
     def __call__(self, x, mask):
         for l in self.layers:
+            l = checkpoint(l) if self.checkpoint else l
             x = l(x, mask)
         return self.ln(x)
 
@@ -191,7 +195,7 @@ class TransformerDecoderLayer(Module):
         mlp_dims: Optional[int] = None,
         dropout: float = 0.0,
         activation: Callable[[Any], Any] = relu,
-        norm_first: bool = False,
+        norm_first: bool = True,
     ):
         super().__init__()
         mlp_dims = mlp_dims or dims * 4
@@ -254,7 +258,8 @@ class TransformerDecoder(Module):
         mlp_dims: Optional[int] = None,
         dropout: float = 0.0,
         activation=relu,
-        norm_first: bool = False,
+        norm_first: bool = True,
+        checkpoint: bool = False,
     ):
         super().__init__()
         self.layers = [
@@ -264,9 +269,11 @@ class TransformerDecoder(Module):
             for i in range(num_layers)
         ]
         self.ln = LayerNorm(dims)
+        self.checkpoint = checkpoint
 
     def __call__(self, x, memory, x_mask, memory_mask):
         for l in self.layers:
+            l = checkpoint(l) if self.checkpoint else l
             x = l(x, memory, x_mask, memory_mask)
         return self.ln(x)
 
@@ -306,7 +313,10 @@ class Transformer(Module):
             standard Transformer decoder. Default: ``None``.
         norm_first (bool, optional): if ``True``, encoder and decoder layers
             will perform layer normalization before attention and MLP
-            operations, otherwise after. Default: ``False``.
+            operations, otherwise after. Default: ``True``.
+        chekpoint (bool, optional): if ``True`` perform gradient checkpointing
+            to reduce the memory usage at the expense of more computation.
+            Default: ``False``.
     """
 
     def __init__(
@@ -320,34 +330,32 @@ class Transformer(Module):
         activation: Callable[[Any], Any] = relu,
         custom_encoder: Optional[Any] = None,
         custom_decoder: Optional[Any] = None,
-        norm_first: bool = False,
+        norm_first: bool = True,
+        checkpoint: bool = False,
     ):
         super().__init__()
-        if custom_encoder is not None:
-            self.encoder = custom_encoder
-        else:
-            self.encoder = TransformerEncoder(
-                num_encoder_layers,
-                dims,
-                num_heads,
-                mlp_dims,
-                dropout,
-                activation,
-                norm_first,
-            )
 
-        if custom_decoder is not None:
-            self.decoder = custom_decoder
-        else:
-            self.decoder = TransformerDecoder(
-                num_decoder_layers,
-                dims,
-                num_heads,
-                mlp_dims,
-                dropout,
-                activation,
-                norm_first,
-            )
+        self.encoder = custom_encoder or TransformerEncoder(
+            num_encoder_layers,
+            dims,
+            num_heads,
+            mlp_dims,
+            dropout,
+            activation,
+            norm_first,
+            checkpoint,
+        )
+
+        self.decoder = custom_decoder or TransformerDecoder(
+            num_decoder_layers,
+            dims,
+            num_heads,
+            mlp_dims,
+            dropout,
+            activation,
+            norm_first,
+            checkpoint,
+        )
 
     def __call__(self, src, tgt, src_mask, tgt_mask, memory_mask):
         memory = self.encoder(src, src_mask)

python/mlx/nn/losses.py

@@ -31,9 +31,14 @@ def cross_entropy(
     Computes the cross entropy loss.
 
     Args:
-        logits (array): The unnormalized predicted logits.
-        targets (array): The target values, as class indices.
-        weights (array, optional): Weights for each target. Default: ``None``.
+        logits (array): The unnormalized logits.
+        targets (array): The ground truth values. These can be class indices or
+            probabilities for each class. If the ``targets`` are class indices,
+            then ``targets`` shape should match the ``logits`` shape with
+            the ``axis`` dimension removed. If the ``targets`` are probabilities
+            (or one-hot encoded), then the ``targets`` shape should be the same as
+            the ``logits`` shape.
+        weights (array, optional): Optional weights for each target. Default: ``None``.
         axis (int, optional): The axis over which to compute softmax. Default: ``-1``.
         label_smoothing (float, optional): Label smoothing factor. Default: ``0``.
         reduction (str, optional): Specifies the reduction to apply to the output:
@@ -41,11 +46,47 @@ def cross_entropy(
 
     Returns:
         array: The computed cross entropy loss.
+
+    Examples:
+        >>> import mlx.core as mx
+        >>> import mlx.nn as nn
+        >>>
+        >>> # Class indices as targets
+        >>> logits = mx.array([[2.0, -1.0], [-1.0, 2.0]])
+        >>> targets = mx.array([0, 1])
+        >>> nn.losses.cross_entropy(logits, targets)
+        array([0.0485873, 0.0485873], dtype=float32)
+        >>>
+        >>> # Probabilities (or one-hot vectors) as targets
+        >>> logits = mx.array([[2.0, -1.0], [-1.0, 2.0]])
+        >>> targets = mx.array([[0.9, 0.1], [0.1, 0.9]])
+        >>> nn.losses.cross_entropy(logits, targets)
+        array([0.348587, 0.348587], dtype=float32)
     """
     if label_smoothing < 0 or label_smoothing >= 1:
         raise ValueError(f"Label smoothing must in [0, 1), got {label_smoothing}.")
 
-    score = mx.take_along_axis(logits, targets[..., None], axis).squeeze(-1)
+    # Whether targets are class indices or probabilities
+    targets_as_probs = targets.ndim == logits.ndim
+
+    def _drop_dim(shape, axis):
+        shape = list(shape)
+        shape.pop(axis)
+        return tuple(shape)
+
+    # Check shapes in two cases: targets as class indices and targets as probabilities
+    if (targets_as_probs and targets.shape != logits.shape) or (
+        not targets_as_probs and targets.shape != _drop_dim(logits.shape, axis)
+    ):
+        raise ValueError(
+            f"Targets shape {targets.shape} does not match logits shape {logits.shape}."
+        )
+
+    if targets_as_probs:
+        score = mx.sum(logits * targets, axis=axis)
+    else:
+        score = mx.take_along_axis(logits, targets[..., None], axis).squeeze(-1)
+
     logsumexp_logits = mx.logsumexp(logits, axis=axis)
     if label_smoothing > 0:
         # Adjust the true class score with label smoothing
@@ -62,10 +103,10 @@ def cross_entropy(
 
     # Apply weights if provided
     if weights is not None:
-        if weights.shape != targets.shape:
+        if weights.shape != loss.shape:
             raise ValueError(
                 f"Weights with shape {weights.shape} is not the same as "
-                f"targets with shape {targets.shape}."
+                f"output loss with shape {loss.shape}."
             )
         loss *= weights
 
@@ -74,29 +115,50 @@ def cross_entropy(
 
 
 def binary_cross_entropy(
-    logits: mx.array, targets: mx.array, reduction: Reduction = "none"
+    inputs: mx.array,
+    targets: mx.array,
+    with_logits: bool = True,
+    reduction: Reduction = "mean",
 ) -> mx.array:
     """
     Computes the binary cross entropy loss.
 
     Args:
-        logits (array): The unnormalized (pre-sigmoid) predicted logits.
+        inputs (array): The predicted values. If ``with_logits`` is ``True``, then
+            ``inputs`` are unnormalized logits. Otherwise, ``inputs`` are probabilities.
         targets (array): The binary target values in {0, 1}.
+        with_logits (bool, optional): Whether ``inputs`` are logits. Default: ``True``.
         reduction (str, optional): Specifies the reduction to apply to the output:
-          ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'none'``.
+          ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'mean'``.
 
     Returns:
         array: The computed binary cross entropy loss.
     Examples:
         >>> import mlx.core as mx
         >>> import mlx.nn as nn
-        >>> inputs = mx.array([0.105361, 0.223144, 1.20397, 0.916291])
+
+        >>> logits = mx.array([0.105361, 0.223144, 1.20397, 0.916291])
         >>> targets = mx.array([0, 0, 1, 1])
-        >>> loss = nn.losses.binary_cross_entropy(inputs, targets, "mean")
+        >>> loss = nn.losses.binary_cross_entropy(logits, targets, reduction="mean")
         >>> loss
-        array([0.612192], dtype=float32)
+        array(0.539245, dtype=float32)
+
+        >>> probs = mx.array([0.1, 0.1, 0.4, 0.4])
+        >>> targets = mx.array([0, 0, 1, 1])
+        >>> loss = nn.losses.binary_cross_entropy(probs, targets, with_logits=False, reduction="mean")
+        >>> loss
+        array(0.510826, dtype=float32)
     """
-    loss = mx.logaddexp(0.0, logits) - targets * logits
+    if inputs.shape != targets.shape:
+        raise ValueError(
+            f"Inputs shape {inputs.shape} does not match targets shape {targets.shape}."
+        )
+
+    if with_logits:
+        loss = mx.logaddexp(0.0, inputs) - inputs * targets
+    else:
+        loss = -(targets * mx.log(inputs) + (1 - targets) * mx.log(1 - inputs))
+
     return _reduce(loss, reduction)
 
 

python/mlx/nn/utils.py

@@ -1,11 +1,14 @@
-# Copyright © 2023 Apple Inc.
+# Copyright © 2023-2024 Apple Inc.
 
+from functools import wraps
 from typing import Callable
 
 import mlx.core as mx
 
+from .layers.base import Module
 
-def value_and_grad(model: "mlx.nn.Module", fn: Callable):
+
+def value_and_grad(model: Module, fn: Callable):
     """Transform the passed function ``fn`` to a function that computes the
     gradients of ``fn`` wrt the model's trainable parameters and also its
     value.
@@ -26,8 +29,42 @@ def value_and_grad(model: "mlx.nn.Module", fn: Callable):
 
     value_grad_fn = mx.value_and_grad(inner_fn)
 
+    @wraps(fn)
     def wrapped_value_grad_fn(*args, **kwargs):
         value, grad = value_grad_fn(model.trainable_parameters(), *args, **kwargs)
         return value, grad
 
     return wrapped_value_grad_fn
+
+
+def checkpoint(module: Module, fn: Callable = None):
+    """Transform the passed callable to one that performs gradient
+    checkpointing with respect to the trainable parameters of the module (and
+    the callable's inputs).
+
+    Args:
+        module (mlx.nn.Module): The module for whose parameters we will be
+            performing gradient checkpointing.
+        fn (Callable, optional): The function to checkpoint. If not provided it
+            defaults to the provided module.
+
+    Returns:
+        A callable that saves the inputs and outputs during the forward pass
+        and recomputes all intermediate states during the backward pass.
+    """
+    if fn is None:
+        # Capturing module instead of module.__call__ allows someone to
+        # monkey-patch __call__ later on and the correct method will be used
+        fn = module
+
+    def inner_fn(params, *args, **kwargs):
+        module.update(params)
+        return fn(*args, **kwargs)
+
+    checkpointed_fn = mx.checkpoint(inner_fn)
+
+    @wraps(fn)
+    def wrapped_checkpointed_fn(*args, **kwargs):
+        return checkpointed_fn(module.trainable_parameters(), *args, **kwargs)
+
+    return wrapped_checkpointed_fn

python/mlx/optimizers.py

@@ -1,7 +1,7 @@
 # Copyright © 2023 Apple Inc.
 
 import math
-from typing import List
+from typing import List, Optional, Tuple
 
 import mlx.core as mx
 from mlx.utils import tree_map
@@ -76,7 +76,7 @@ class Optimizer:
 
 
 class SGD(Optimizer):
-    r"""Stochastic gradient descent optimizer.
+    r"""The stochastic gradient descent optimizer.
 
     Updates a parameter :math:`w` with a gradient :math:`g` as follows
 
@@ -118,18 +118,21 @@ class SGD(Optimizer):
     ):
         """Performs the SGD parameter update and stores :math:`v` in the
         optimizer state."""
-        if self.momentum <= 0:
-            return parameter - self.learning_rate * gradient
-
-        v = state.get("v", mx.zeros_like(gradient))
 
         if self.weight_decay != 0:
             gradient += self.weight_decay * parameter
 
-        v = self.momentum * v
+        if self.momentum <= 0:
+            return parameter - self.learning_rate * gradient
+
         if self.dampening > 0:
+            v = (
+                state.get("v", (self.dampening / self.momentum) * gradient)
+                * self.momentum
+            )
             v += (1 - self.dampening) * gradient
         else:
+            v = state.get("v", mx.zeros_like(gradient)) * self.momentum
             v += gradient
 
         if self.nesterov:
@@ -141,7 +144,7 @@ class SGD(Optimizer):
 
 
 class RMSprop(Optimizer):
-    r"""Implementation of the RMSprop optimizer [1].
+    r"""The RMSprop optimizer [1].
 
     [1]: Tieleman, T. and Hinton, G. 2012. Lecture 6.5-rmsprop, coursera: Neural networks for machine learning
 
@@ -190,7 +193,7 @@ class RMSprop(Optimizer):
 
 
 class Adagrad(Optimizer):
-    r"""Implementation of the Adagrad optimizer [1].
+    r"""The Adagrad optimizer [1].
 
     Our Adagrad implementation follows the original paper. In detail,
 
@@ -235,7 +238,7 @@ class Adagrad(Optimizer):
 
 
 class AdaDelta(Optimizer):
-    r"""Implementation of the AdaDelta optimizer with learning rate[1].
+    r"""The AdaDelta optimizer with a learning rate [1].
 
     Our AdaDelta implementation follows the original paper. In detail,
 
@@ -281,7 +284,7 @@ class AdaDelta(Optimizer):
         eps = self.eps
 
         v = state.get("v", mx.zeros_like(gradient))
-        u = state.get("s", mx.zeros_like(gradient))
+        u = state.get("u", mx.zeros_like(gradient))
 
         v = rho * v + (1 - rho) * mx.square(gradient)
         d = mx.sqrt(u + eps) / mx.sqrt(v + eps) * gradient
@@ -294,7 +297,7 @@ class AdaDelta(Optimizer):
 
 
 class Adam(Optimizer):
-    r"""Implementation of the Adam optimizer [1].
+    r"""The Adam optimizer [1].
 
     Our Adam implementation follows the original paper and omits the bias
     correction in the first and second moment estimates. In detail,
@@ -346,7 +349,7 @@ class Adam(Optimizer):
 
 
 class AdamW(Adam):
-    r"""Implementation of the AdamW optimizer [1].
+    r"""The AdamW optimizer [1].
 
     Following the above convention, in contrast with [1], we do not use bias
     correction in the first and second moments for AdamW. We update the weights
@@ -395,8 +398,7 @@ class AdamW(Adam):
 
 
 class Adamax(Adam):
-    r"""Implementation of the Adamax optimizer. It is a variant of Adam based
-    on the infinity norm [1].
+    r"""The Adamax optimizer, a variant of Adam based on the infinity norm [1].
 
     Our Adam implementation follows the original paper and omits the bias
     correction in the first and second moment estimates. In detail,
@@ -423,6 +425,10 @@ class Adamax(Adam):
         self, learning_rate: float, betas: List[float] = [0.9, 0.999], eps: float = 1e-8
     ):
         super().__init__(learning_rate, betas, eps)
+        if not 0.0 <= eps:
+            raise ValueError(
+                f"Epsilon value should be >=0, {self.eps} was provided instead"
+            )
 
     def apply_single(
         self, gradient: mx.array, parameter: mx.array, state: OptimizerState
@@ -445,7 +451,7 @@ class Adamax(Adam):
 
 
 class Lion(Optimizer):
-    r"""Implementation of the Lion optimizer [1].
+    r"""The Lion optimizer [1].
 
     Since updates are computed through the sign operation, they tend to
     have larger norm than for other optimizers such as SGD and Adam.
@@ -459,8 +465,8 @@ class Lion(Optimizer):
 
     .. math::
 
-        c_{t + 1} &= \beta_1 m_t + (1 - \beta_1) g_t
-        m_{t + 1} &= \beta_2 m_t + (1 - \beta_2) g_t
+        c_{t + 1} &= \beta_1 m_t + (1 - \beta_1) g_t \\
+        m_{t + 1} &= \beta_2 m_t + (1 - \beta_2) g_t \\
         w_{t + 1} &= w_t - \eta (\text{sign}(c_t) + \lambda w_t)
 
     Args:
@@ -498,3 +504,139 @@ class Lion(Optimizer):
         if weight_decay > 0:
             parameter = (1 - lr * weight_decay) * parameter
         return parameter - lr * mx.sign(c)
+
+
+class Adafactor(Optimizer):
+    r"""The Adafactor optimizer.
+
+    Our Adafactor implementation follows the original paper: `Adafactor:
+    Adaptive Learning Rates with Sublinear Memory Cost
+    <https://arxiv.org/abs/1804.04235>`_
+
+    Args:
+        learning_rate (float, optional): The learning rate. Default: ``None``.
+        eps (tuple(float, float), optional): The first term :math:`\epsilon_1`
+            added to the square of the gradients to improve numerical
+            stability and the second term :math:`\epsilon_2` is used for
+            parameter scaling if ``parameter_scale`` is set to ``True``.
+            Default: ``(1e-30, 1e-3)``.
+        clip_threshold (float, optional): Clips the unscaled update at
+            ``clip_threshold``. Default: ``1.0``.
+        decay_rate (float, optional): Coefficient for the running average
+            of the squared gradient. Default: ``-0.8``.
+        beta_1 (float, optional): If set to a value bigger than zero
+            then first moment will be used. Default: ``None``.
+        weight_decay (float, optional): The weight decay :math:`\lambda`.
+            Default: ``0.0``.
+        scale_parameter (bool, optional): If set to ``True`` the learning rate
+            will be scaled by :math:`\max(\epsilon_1, \text{RMS}(w_{t-1}))`.
+            Default: ``True``.
+        relative_step (bool, optional): If set to ``True`` the ``learning_rate``
+            will be ignored and relative step size will be computed.
+            Default: ``True``.
+        warmup_init (bool, optional): If set to ``True`` then the relative
+            step size will be calculated by the current step. Default:
+            ``False``.
+    """
+
+    def __init__(
+        self,
+        learning_rate: Optional[float] = None,
+        eps: Tuple[float, float] = (1e-30, 1e-3),
+        clip_threshold: float = 1.0,
+        decay_rate: float = -0.8,
+        beta_1: Optional[float] = None,
+        weight_decay: float = 0.0,
+        scale_parameter: bool = True,
+        relative_step: bool = True,
+        warmup_init: bool = False,
+    ):
+        super().__init__()
+        self.learning_rate = learning_rate
+        self.eps = eps
+        self.clip_threshold = clip_threshold
+        self.decay_rate = decay_rate
+        self.beta_1 = beta_1
+        self.weight_decay = weight_decay
+        self.scale_parameter = scale_parameter
+        self.relative_step = relative_step
+        self.warmup_init = warmup_init
+
+    def _compute_rms(self, inputs):
+        return mx.sqrt(mx.mean(mx.square(inputs)))
+
+    def _compute_learning_rate(self, step, parameter_rms):
+        relative_step_size = self.learning_rate
+        if self.relative_step:
+            min_step = 1e-6 * step if self.warmup_init else 1e-2
+            relative_step_size = min(min_step, 1 / math.sqrt(step))
+
+        parameter_scale = 1.0
+        if self.scale_parameter:
+            parameter_scale = mx.maximum(self.eps[1], parameter_rms)
+        return parameter_scale * relative_step_size
+
+    def _approximate_exp_moving_avg(self, exp_avg_sq_row, exp_avg_sq_col):
+        r_factor = mx.rsqrt(
+            exp_avg_sq_row / mx.mean(exp_avg_sq_row, axis=-1, keepdims=True)
+        )
+        c_factor = mx.rsqrt(exp_avg_sq_col)
+        return mx.matmul(
+            mx.expand_dims(r_factor, axis=-1), mx.expand_dims(c_factor, axis=0)
+        )
+
+    def apply_single(
+        self, gradient: mx.array, parameter: mx.array, state: OptimizerState
+    ):
+        """Performs the Adafactor parameter and state update."""
+        gradient_shape = gradient.shape
+        factored = len(gradient_shape) >= 2
+        step = state.get("step", 0) + 1
+        state["step"] = step
+        use_first_moment = self.beta_1 is not None
+
+        parameter_rms = self._compute_rms(parameter)
+        learning_rate = self._compute_learning_rate(step, parameter_rms)
+        beta_2 = 1.0 - math.pow(step, self.decay_rate)
+        update = mx.square(gradient) + self.eps[0]
+
+        if factored:
+            exp_avg_sq_row = state.get(
+                "exp_avg_sq_row", mx.zeros(gradient_shape[:-1], dtype=gradient.dtype)
+            )
+            exp_avg_sq_col = state.get(
+                "exp_avg_sq_col",
+                mx.zeros(
+                    gradient_shape[:-2] + gradient_shape[-1:], dtype=gradient.dtype
+                ),
+            )
+            exp_avg_sq_row = (beta_2 * exp_avg_sq_row) + (
+                (1 - beta_2) * mx.mean(update, axis=-1)
+            )
+            exp_avg_sq_col = (beta_2 * exp_avg_sq_col) + (
+                (1 - beta_2) * mx.mean(update, axis=-2)
+            )
+            state["exp_avg_sq_row"] = exp_avg_sq_row
+            state["exp_avg_sq_col"] = exp_avg_sq_col
+            update = self._approximate_exp_moving_avg(exp_avg_sq_row, exp_avg_sq_col)
+            update = update * gradient
+        else:
+            exp_avg_sq = state.get("exp_avg_sq", mx.zeros_like(gradient))
+            exp_avg_sq = (beta_2 * exp_avg_sq) + ((1 - beta_2) * update)
+            state["exp_avg_sq"] = exp_avg_sq
+            update = mx.rsqrt(exp_avg_sq) * gradient
+
+        update = update / mx.maximum(
+            1.0, self._compute_rms(update) / self.clip_threshold
+        )
+        update = learning_rate * update
+
+        if use_first_moment:
+            exp_avg = state.get("exp_avg", mx.zeros_like(gradient))
+            exp_avg = (self.beta_1 * exp_avg) + ((1 - self.beta_1) * update)
+            state["exp_avg"] = exp_avg
+            update = exp_avg
+
+        if self.weight_decay != 0:
+            parameter += parameter * (-self.weight_decay * learning_rate)
+        return parameter - update

python/mlx/py.typed

@@ -0,0 +1 @@
+

python/src/array.cpp

@@ -39,6 +39,10 @@ py::list to_list(array& a, size_t index, int dim) {
 }
 
 auto to_scalar(array& a) {
+  {
+    py::gil_scoped_release nogil;
+    a.eval();
+  }
   switch (a.dtype()) {
     case bool_:
       return py::cast(a.item<bool>());
@@ -73,7 +77,10 @@ py::object tolist(array& a) {
   if (a.ndim() == 0) {
     return to_scalar(a);
   }
-  a.eval();
+  {
+    py::gil_scoped_release nogil;
+    a.eval();
+  }
   py::object pl;
   switch (a.dtype()) {
     case bool_:
@@ -229,9 +236,28 @@ array array_from_list(
       return array(vals.begin(), shape, specified_type.value_or(bool_));
     }
     case pyint: {
-      std::vector<int> vals;
-      fill_vector(pl, vals);
-      return array(vals.begin(), shape, specified_type.value_or(int32));
+      auto dtype = specified_type.value_or(int32);
+      if (dtype == int64) {
+        std::vector<int64_t> vals;
+        fill_vector(pl, vals);
+        return array(vals.begin(), shape, dtype);
+      } else if (dtype == uint64) {
+        std::vector<uint64_t> vals;
+        fill_vector(pl, vals);
+        return array(vals.begin(), shape, dtype);
+      } else if (dtype == uint32) {
+        std::vector<uint32_t> vals;
+        fill_vector(pl, vals);
+        return array(vals.begin(), shape, dtype);
+      } else if (is_floating_point(dtype)) {
+        std::vector<float> vals;
+        fill_vector(pl, vals);
+        return array(vals.begin(), shape, dtype);
+      } else {
+        std::vector<int> vals;
+        fill_vector(pl, vals);
+        return array(vals.begin(), shape, dtype);
+      }
     }
     case pyfloat: {
       std::vector<float> vals;
@@ -625,6 +651,7 @@ void init_array(py::module_& m) {
       .def_buffer([](array& a) {
         // Eval if not already evaled
         if (!a.is_evaled()) {
+          py::gil_scoped_release nogil;
           a.eval();
         }
         return pybind11::buffer_info(
@@ -648,17 +675,14 @@ void init_array(py::module_& m) {
           "nbytes",
           &array::nbytes,
           R"pbdoc(The number of bytes in the array.)pbdoc")
-      // TODO, this makes a deep copy of the shape
-      // implement alternatives to use reference
-      // https://pybind11.readthedocs.io/en/stable/advanced/cast/stl.html
       .def_property_readonly(
           "shape",
-          [](const array& a) { return a.shape(); },
+          [](const array& a) { return py::tuple(py::cast(a.shape())); },
           R"pbdoc(
           The shape of the array as a Python list.
 
           Returns:
-            list(int): A list containing the sizes of each dimension.
+            tuple(int): A tuple containing the sizes of each dimension.
         )pbdoc")
       .def_property_readonly(
           "dtype",
@@ -923,6 +947,7 @@ void init_array(py::module_& m) {
           "__repr__",
           [](array& a) {
             if (!a.is_evaled()) {
+              py::gil_scoped_release nogil;
               a.eval();
             }
             std::ostringstream os;
@@ -1458,5 +1483,26 @@ void init_array(py::module_& m) {
           "decimals"_a = 0,
           py::kw_only(),
           "stream"_a = none,
-          "See :func:`round`.");
+          "See :func:`round`.")
+      .def(
+          "diagonal",
+          [](const array& a,
+             int offset,
+             int axis1,
+             int axis2,
+             StreamOrDevice s) { return diagonal(a, offset, axis1, axis2, s); },
+          "offset"_a = 0,
+          "axis1"_a = 0,
+          "axis2"_a = 1,
+          "stream"_a = none,
+          "See :func:`diagonal`.")
+      .def(
+          "diag",
+          [](const array& a, int k, StreamOrDevice s) { return diag(a, k, s); },
+          "k"_a = 0,
+          py::kw_only(),
+          "stream"_a = none,
+          R"pbdoc(
+            Extract a diagonal or construct a diagonal matrix.
+        )pbdoc");
 }

python/src/linalg.cpp

@@ -177,4 +177,37 @@ void init_linalg(py::module_& parent_module) {
           >>> la.norm(m[0, :, :]), LA.norm(m[1, :, :])
           (array(3.74166, dtype=float32), array(11.225, dtype=float32))
       )pbdoc");
+  m.def(
+      "qr",
+      &qr,
+      "a"_a,
+      py::kw_only(),
+      "stream"_a = none,
+      R"pbdoc(
+        qr(a: array, *, stream: Union[None, Stream, Device] = None) -> (array, array)
+
+        The QR factorizatoin of the input matrix.
+
+        This function supports arrays with at least 2 dimensions. The matrices
+        which are factorized are assumed to be in the last two dimensions of
+        the input.
+
+        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 ``Q`` and ``R`` matrices.
+
+        Example:
+            >>> A = mx.array([[2., 3.], [1., 2.]])
+            >>> Q, R = mx.linalg.qr(A, stream=mx.cpu)
+            >>> Q
+            array([[-0.894427, -0.447214],
+                   [-0.447214, 0.894427]], dtype=float32)
+            >>> R
+            array([[-2.23607, -3.57771],
+                   [0, 0.447214]], dtype=float32)
+      )pbdoc");
 }

python/src/load.cpp

@@ -181,9 +181,10 @@ std::unordered_map<std::string, array> mlx_load_safetensor_helper(
       "[load_safetensors] Input must be a file-like object, or string");
 }
 
-std::unordered_map<std::string, array> mlx_load_gguf_helper(
-    py::object file,
-    StreamOrDevice s) {
+std::pair<
+    std::unordered_map<std::string, array>,
+    std::unordered_map<std::string, MetaData>>
+mlx_load_gguf_helper(py::object file, StreamOrDevice s) {
   if (py::isinstance<py::str>(file)) { // Assume .gguf file path string
     return load_gguf(py::cast<std::string>(file), s);
   }
@@ -194,6 +195,8 @@ std::unordered_map<std::string, array> mlx_load_gguf_helper(
 std::unordered_map<std::string, array> mlx_load_npz_helper(
     py::object file,
     StreamOrDevice s) {
+  bool own_file = py::isinstance<py::str>(file);
+
   py::module_ zipfile = py::module_::import("zipfile");
   if (!is_zip_file(zipfile, file)) {
     throw std::invalid_argument(
@@ -222,9 +225,11 @@ std::unordered_map<std::string, array> mlx_load_npz_helper(
   }
 
   // If we don't own the stream and it was passed to us, eval immediately
-  for (auto& [key, arr] : array_dict) {
+  if (!own_file) {
     py::gil_scoped_release gil;
-    arr.eval();
+    for (auto& [key, arr] : array_dict) {
+      arr.eval();
+    }
   }
 
   return array_dict;
@@ -246,9 +251,10 @@ array mlx_load_npy_helper(py::object file, StreamOrDevice s) {
       "[load_npy] Input must be a file-like object, or string");
 }
 
-DictOrArray mlx_load_helper(
+LoadOutputTypes mlx_load_helper(
     py::object file,
     std::optional<std::string> format,
+    bool return_metadata,
     StreamOrDevice s) {
   if (!format.has_value()) {
     std::string fname;
@@ -258,7 +264,7 @@ DictOrArray mlx_load_helper(
       fname = file.attr("name").cast<std::string>();
     } else {
       throw std::invalid_argument(
-          "[load] Input must be a file-like object, or string");
+          "[load] Input must be a file-like object opened in binary mode, or string");
     }
     size_t ext = fname.find_last_of('.');
     if (ext == std::string::npos) {
@@ -268,6 +274,10 @@ DictOrArray mlx_load_helper(
     format.emplace(fname.substr(ext + 1));
   }
 
+  if (return_metadata && format.value() != "gguf") {
+    throw std::invalid_argument(
+        "[load] metadata not supported for format " + format.value());
+  }
   if (format.value() == "safetensors") {
     return mlx_load_safetensor_helper(file, s);
   } else if (format.value() == "npz") {
@@ -275,7 +285,12 @@ DictOrArray mlx_load_helper(
   } else if (format.value() == "npy") {
     return mlx_load_npy_helper(file, s);
   } else if (format.value() == "gguf") {
-    return mlx_load_gguf_helper(file, s);
+    auto [weights, metadata] = mlx_load_gguf_helper(file, s);
+    if (return_metadata) {
+      return std::make_pair(weights, metadata);
+    } else {
+      return weights;
+    }
   } else {
     throw std::invalid_argument("[load] Unknown file format " + format.value());
   }
@@ -421,7 +436,7 @@ void mlx_savez_helper(
     auto py_ostream = zipfile_object.open(fname, 'w');
     auto writer = std::make_shared<PyFileWriter>(py_ostream);
     {
-      py::gil_scoped_release gil;
+      py::gil_scoped_release nogil;
       save(writer, a);
     }
   }
@@ -432,28 +447,42 @@ void mlx_savez_helper(
 void mlx_save_safetensor_helper(py::object file, py::dict d) {
   auto arrays_map = d.cast<std::unordered_map<std::string, array>>();
   if (py::isinstance<py::str>(file)) {
-    save_safetensors(py::cast<std::string>(file), arrays_map);
-    return;
+    {
+      py::gil_scoped_release nogil;
+      save_safetensors(py::cast<std::string>(file), arrays_map);
+    }
   } else if (is_ostream_object(file)) {
     auto writer = std::make_shared<PyFileWriter>(file);
     {
-      py::gil_scoped_release gil;
+      py::gil_scoped_release nogil;
       save_safetensors(writer, arrays_map);
     }
-
-    return;
+  } else {
+    throw std::invalid_argument(
+        "[save_safetensors] Input must be a file-like object, or string");
   }
-
-  throw std::invalid_argument(
-      "[save_safetensors] Input must be a file-like object, or string");
 }
 
-void mlx_save_gguf_helper(py::object file, py::dict d) {
-  auto arrays_map = d.cast<std::unordered_map<std::string, array>>();
+void mlx_save_gguf_helper(
+    py::object file,
+    py::dict a,
+    std::optional<py::dict> m) {
+  auto arrays_map = a.cast<std::unordered_map<std::string, array>>();
   if (py::isinstance<py::str>(file)) {
-    save_gguf(py::cast<std::string>(file), arrays_map);
-    return;
+    if (m) {
+      auto metadata_map =
+          m.value().cast<std::unordered_map<std::string, MetaData>>();
+      {
+        py::gil_scoped_release nogil;
+        save_gguf(py::cast<std::string>(file), arrays_map, metadata_map);
+      }
+    } else {
+      {
+        py::gil_scoped_release nogil;
+        save_gguf(py::cast<std::string>(file), arrays_map);
+      }
+    }
+  } else {
+    throw std::invalid_argument("[save_gguf] Input must be a string");
   }
-
-  throw std::invalid_argument("[save_safetensors] Input must be a string");
 }

python/src/load.h

@@ -7,26 +7,36 @@
 #include <string>
 #include <unordered_map>
 #include <variant>
-#include "mlx/ops.h"
+#include "mlx/io.h"
 
 namespace py = pybind11;
 using namespace mlx::core;
 
-using DictOrArray = std::variant<array, std::unordered_map<std::string, array>>;
+using LoadOutputTypes = std::variant<
+    array,
+    std::unordered_map<std::string, array>,
+    std::pair<
+        std::unordered_map<std::string, array>,
+        std::unordered_map<std::string, MetaData>>>;
 
 std::unordered_map<std::string, array> mlx_load_safetensor_helper(
     py::object file,
     StreamOrDevice s);
 void mlx_save_safetensor_helper(py::object file, py::dict d);
 
-std::unordered_map<std::string, array> mlx_load_gguf_helper(
+std::pair<
+    std::unordered_map<std::string, array>,
+    std::unordered_map<std::string, MetaData>>
+mlx_load_gguf_helper(py::object file, StreamOrDevice s);
+void mlx_save_gguf_helper(
     py::object file,
-    StreamOrDevice s);
-void mlx_save_gguf_helper(py::object file, py::dict d);
+    py::dict d,
+    std::optional<py::dict> m);
 
-DictOrArray mlx_load_helper(
+LoadOutputTypes mlx_load_helper(
     py::object file,
     std::optional<std::string> format,
+    bool return_metadata,
     StreamOrDevice s);
 void mlx_save_helper(py::object file, array a);
 void mlx_savez_helper(

python/src/ops.cpp

@@ -55,7 +55,7 @@ void init_ops(py::module_& m) {
         Args:
             a (array): Input array.
             shape (tuple(int)): New shape.
-            stream (Stream, optional): Stream or device. Defaults to ```None```
+            stream (Stream, optional): Stream or device. Defaults to ``None``
               in which case the default stream of the default device is used.
 
         Returns:
@@ -78,6 +78,11 @@ void init_ops(py::module_& m) {
 
       Flatten an array.
 
+      The axes flattened will be between ``start_axis`` and ``end_axis``,
+      inclusive. Negative axes are supported. After converting negative axis to
+      positive, axes outside the valid range will be clamped to a valid value,
+      ``start_axis`` to ``0`` and ``end_axis`` to ``ndim - 1``.
+
       Args:
           a (array): Input array.
           start_axis (int, optional): The first dimension to flatten. Defaults to ``0``.
@@ -87,6 +92,14 @@ void init_ops(py::module_& m) {
 
       Returns:
           array: The flattened array.
+
+      Example:
+          >>> a = mx.array([[1, 2], [3, 4]])
+          >>> mx.flatten(a)
+          array([1, 2, 3, 4], dtype=int32)
+          >>>
+          >>> mx.flatten(a, start_axis=0, end_axis=-1)
+          array([1, 2, 3, 4], dtype=int32)
   )pbdoc");
   m.def(
       "squeeze",
@@ -112,7 +125,7 @@ void init_ops(py::module_& m) {
         Args:
             a (array): Input array.
             axis (int or tuple(int), optional): Axes to remove. Defaults
-            to ```None``` in which case all size one axes are removed.
+            to ``None`` in which case all size one axes are removed.
 
         Returns:
             array: The output array with size one axes removed.
@@ -566,7 +579,7 @@ void init_ops(py::module_& m) {
         Args:
             a (array): Input array or scalar.
             b (array): Input array or scalar.
-            equal_nan (bool): If ``True``, NaNs are treated as equal.
+            equal_nan (bool): If ``True``, NaNs are considered equal.
               Defaults to ``False``.
 
         Returns:
@@ -801,7 +814,7 @@ void init_ops(py::module_& m) {
         Element-wise error function.
 
         .. math::
-          \mathrm{erf}(x) = \frac{2}{\sqrt{\pi}} \int_0^t e^{-t^2} \, dx
+          \mathrm{erf}(x) = \frac{2}{\sqrt{\pi}} \int_0^x e^{-t^2} \, dt
 
         Args:
             a (array): Input array.
@@ -1648,12 +1661,15 @@ void init_ops(py::module_& m) {
       "rtol"_a = 1e-5,
       "atol"_a = 1e-8,
       py::kw_only(),
+      "equal_nan"_a = false,
       "stream"_a = none,
       R"pbdoc(
-        allclose(a: array, b: array, /, rtol: float = 1e-05, atol: float = 1e-08, *, stream: Union[None, Stream, Device] = None) -> array
+        allclose(a: array, b: array, /, rtol: float = 1e-05, atol: float = 1e-08, *, equal_nan: bool = False, stream: Union[None, Stream, Device] = None) -> array
 
         Approximate comparison of two arrays.
 
+        Infinite values are considered equal if they have the same sign, NaN values are not equal unless ``equal_nan`` is ``True``.
+
         The arrays are considered equal if:
 
         .. code-block::
@@ -1668,6 +1684,47 @@ void init_ops(py::module_& m) {
             b (array): Input array.
             rtol (float): Relative tolerance.
             atol (float): Absolute tolerance.
+            equal_nan (bool): If ``True``, NaNs are considered equal.
+              Defaults to ``False``.
+
+        Returns:
+            array: The boolean output scalar indicating if the arrays are close.
+      )pbdoc");
+  m.def(
+      "isclose",
+      &isclose,
+      "a"_a,
+      "b"_a,
+      py::pos_only(),
+      "rtol"_a = 1e-5,
+      "atol"_a = 1e-8,
+      py::kw_only(),
+      "equal_nan"_a = false,
+      "stream"_a = none,
+      R"pbdoc(
+        isclose(a: array, b: array, /, rtol: float = 1e-05, atol: float = 1e-08, *, equal_nan: bool = False, stream: Union[None, Stream, Device] = None) -> array
+
+        Returns a boolean array where two arrays are element-wise equal within a tolerance.
+
+        Infinite values are considered equal if they have the same sign, NaN values are
+        not equal unless ``equal_nan`` is ``True``.
+
+        Two values are considered equal if:
+
+        .. code-block::
+
+         abs(a - b) <= (atol + rtol * abs(b))
+
+        Note unlike :func:`array_equal`, this function supports numpy-style
+        broadcasting.
+
+        Args:
+            a (array): Input array.
+            b (array): Input array.
+            rtol (float): Relative tolerance.
+            atol (float): Absolute tolerance.
+            equal_nan (bool): If ``True``, NaNs are considered equal.
+              Defaults to ``False``.
 
         Returns:
             array: The boolean output scalar indicating if the arrays are close.
@@ -1867,11 +1924,11 @@ void init_ops(py::module_& m) {
         isposinf(a: array, stream: Union[None, Stream, Device] = None) -> array
 
         Return a boolean array indicating which elements are positive infinity.
-        
+
         Args:
             a (array): Input array.
             stream (Union[None, Stream, Device]): Optional stream or device.
-        
+
         Returns:
             array: The boolean array indicating which elements are positive infinity.
       )pbdoc");
@@ -1886,11 +1943,11 @@ void init_ops(py::module_& m) {
         isneginf(a: array, stream: Union[None, Stream, Device] = None) -> array
 
         Return a boolean array indicating which elements are negative infinity.
-        
+
         Args:
             a (array): Input array.
             stream (Union[None, Stream, Device]): Optional stream or device.
-        
+
         Returns:
             array: The boolean array indicating which elements are negative infinity.
       )pbdoc");
@@ -2254,7 +2311,7 @@ void init_ops(py::module_& m) {
               singleton dimensions, defaults to `False`.
 
         Returns:
-            array: The output array with the indices of the minimum values.
+            array: The ``uint32`` array with the indices of the minimum values.
       )pbdoc");
   m.def(
       "argmax",
@@ -2287,7 +2344,7 @@ void init_ops(py::module_& m) {
               singleton dimensions, defaults to `False`.
 
         Returns:
-            array: The output array with the indices of the maximum values.
+            array: The ``uint32`` array with the indices of the maximum values.
       )pbdoc");
   m.def(
       "sort",
@@ -2343,7 +2400,7 @@ void init_ops(py::module_& m) {
               If unspecified, it defaults to -1 (sorting over the last axis).
 
         Returns:
-            array: The indices that sort the input array.
+            array: The ``uint32`` array containing indices that sort the input.
       )pbdoc");
   m.def(
       "partition",
@@ -2416,7 +2473,7 @@ void init_ops(py::module_& m) {
               If unspecified, it defaults to ``-1``.
 
         Returns:
-            array: The indices that partition the input array.
+            array: The `uint32`` array containing indices that partition the input.
       )pbdoc");
   m.def(
       "topk",
@@ -2886,6 +2943,10 @@ void init_ops(py::module_& m) {
           throw std::invalid_argument("[convolve] Inputs must be 1D.");
         }
 
+        if (a.size() == 0 || v.size() == 0) {
+          throw std::invalid_argument("[convolve] Inputs cannot be empty.");
+        }
+
         array in = a.size() < v.size() ? v : a;
         array wt = a.size() < v.size() ? a : v;
         wt = slice(wt, {wt.shape(0) - 1}, {-wt.shape(0) - 1}, {-1}, s);
@@ -3117,10 +3178,11 @@ void init_ops(py::module_& m) {
       "file"_a,
       py::pos_only(),
       "format"_a = none,
+      "return_metadata"_a = false,
       py::kw_only(),
       "stream"_a = none,
       R"pbdoc(
-        load(file: str, /, format: Optional[str] = None, *, stream: Union[None, Stream, Device] = None) -> Union[array, Dict[str, array]]
+        load(file: str, /, format: Optional[str] = None, return_metadata: bool = False, *, stream: Union[None, Stream, Device] = None) -> Union[array, Dict[str, array]]
 
         Load array(s) from a binary file.
 
@@ -3131,10 +3193,15 @@ void init_ops(py::module_& m) {
             format (str, optional): Format of the file. If ``None``, the format
               is inferred from the file extension. Supported formats: ``npy``,
               ``npz``, and ``safetensors``. Default: ``None``.
+            return_metadata (bool, optional): Load the metadata for formats which
+              support matadata. The metadata will be returned as an additional
+              dictionary.
         Returns:
             result (array, dict):
                 A single array if loading from a ``.npy`` file or a dict mapping
                 names to arrays if loading from a ``.npz`` or ``.safetensors`` file.
+                If ``return_metadata` is ``True`` an additional dictionary of metadata
+                will be returned.
 
         Warning:
 
@@ -3164,8 +3231,9 @@ void init_ops(py::module_& m) {
       &mlx_save_gguf_helper,
       "file"_a,
       "arrays"_a,
+      "metadata"_a = none,
       R"pbdoc(
-        save_gguf(file: str, arrays: Dict[str, array])
+        save_gguf(file: str, arrays: Dict[str, array], metadata: Dict[str, Union[array, str, List[str]]])
 
         Save array(s) to a binary file in ``.gguf`` format.
 
@@ -3175,6 +3243,9 @@ void init_ops(py::module_& m) {
         Args:
             file (file, str): File in which the array is saved.
             arrays (dict(str, array)): The dictionary of names to arrays to be saved.
+            metadata (dict(str, Union[array, str, list(str)])): The dictionary of
+               metadata to be saved. The values can be a scalar or 1D obj:`array`,
+               a :obj:`str`, or a :obj:`list` of :obj:`str`.
       )pbdoc");
   m.def(
       "where",
@@ -3390,20 +3461,20 @@ void init_ops(py::module_& m) {
       "a"_a,
       "b"_a,
       py::pos_only(),
-      "dims"_a = 2,
+      "axes"_a = 2,
       py::kw_only(),
       "stream"_a = none,
       R"pbdoc(
-        tensordot(a: array, b: array, /, dims: Union[int, List[List[int]]] = 2, *, stream: Union[None, Stream, Device] = None) -> array
+        tensordot(a: array, b: array, /, axes: Union[int, List[List[int]]] = 2, *, stream: Union[None, Stream, Device] = None) -> array
 
         Compute the tensor dot product along the specified axes.
 
         Args:
           a (array): Input array
           b (array): Input array
-          dims (int or list(list(int)), optional): The number of dimensions to
+          axes (int or list(list(int)), optional): The number of dimensions to
             sum over. If an integer is provided, then sum over the last
-            ``dims`` dimensions of ``a`` and the first ``dims`` dimensions of
+            ``axes`` dimensions of ``a`` and the first ``axes`` dimensions of
             ``b``. If a list of lists is provided, then sum over the
             corresponding dimensions of ``a`` and ``b``. (default: 2)
 
@@ -3499,11 +3570,68 @@ void init_ops(py::module_& m) {
             c (array): Input array or scalar.
             a (array): Input array or scalar.
             b (array): Input array or scalar.
-            alpha (float, optional): Scaling factor for the 
+            alpha (float, optional): Scaling factor for the
                 matrix product of ``a`` and ``b`` (default: ``1``)
             beta (float, optional): Scaling factor for ``c`` (default: ``1``)
 
         Returns:
             array: ``alpha * (a @ b)  + beta * c``
       )pbdoc");
+  m.def(
+      "diagonal",
+      &diagonal,
+      "a"_a,
+      "offset"_a = 0,
+      "axis1"_a = 0,
+      "axis2"_a = 1,
+      "stream"_a = none,
+      R"pbdoc(
+        diagonal(a: array, offset: int = 0, axis1: int = 0, axis2: int = 1, stream: Union[None, Stream, Device] = None) -> array
+
+        Return specified diagonals.
+
+        If ``a`` is 2-D, then a 1-D array containing the diagonal at the given
+        ``offset`` is returned.
+
+        If ``a`` has more than two dimensions, then ``axis1`` and ``axis2``
+        determine the 2D subarrays from which diagonals are extracted. The new
+        shape is the original shape with ``axis1`` and ``axis2`` removed and a
+        new dimension inserted at the end corresponding to the diagonal.
+
+        Args:
+          a (array): Input array
+          offset (int, optional): Offset of the diagonal from the main diagonal.
+            Can be positive or negative. Default: ``0``.
+          axis1 (int, optional): The first axis of the 2-D sub-arrays from which
+              the diagonals should be taken. Default: ``0``.
+          axis2 (int, optional): The second axis of the 2-D sub-arrays from which
+              the diagonals should be taken. Default: ``1``.
+
+        Returns:
+            array: The diagonals of the array.
+      )pbdoc");
+  m.def(
+      "diag",
+      &diag,
+      "a"_a,
+      py::pos_only(),
+      "k"_a = 0,
+      py::kw_only(),
+      "stream"_a = none,
+      R"pbdoc(
+        diag(a: array, /, k: int = 0, *, stream: Union[None, Stream, Device] = None) -> array
+
+        Extract a diagonal or construct a diagonal matrix.
+        If ``a`` is 1-D then a diagonal matrix is constructed with ``a`` on the
+        :math:`k`-th diagonal. If ``a`` is 2-D then the :math:`k`-th diagonal is
+        returned.
+
+        Args:
+            a (array): 1-D or 2-D input array.
+            k (int, optional): The diagonal to extract or construct.
+                Default: ``0``.
+
+        Returns:
+            array: The extracted diagonal or the constructed diagonal matrix.
+        )pbdoc");
 }

python/src/transforms.cpp

@@ -1,6 +1,4 @@
-// Copyright © 2023 Apple Inc.
-
-#include <pybind11/functional.h>
+// Copyright © 2023-2024 Apple Inc.
 #include <pybind11/pybind11.h>
 #include <pybind11/stl.h>
 #include <algorithm>
@@ -143,7 +141,8 @@ std::vector<array> tree_flatten(py::object tree, bool strict = true) {
     if (py::isinstance<array>(obj)) {
       flat_tree.push_back(py::cast<array>(obj));
     } else if (strict) {
-      throw std::invalid_argument("Argument is not an array");
+      throw std::invalid_argument(
+          "[tree_flatten] The argument should contain only arrays");
     }
   });
 
@@ -163,6 +162,55 @@ py::object tree_unflatten(
   });
 }
 
+py::object structure_sentinel() {
+  static py::object sentinel;
+
+  if (sentinel.ptr() == nullptr) {
+    sentinel = py::capsule(&sentinel);
+    // probably not needed but this should make certain that we won't ever
+    // delete the sentinel
+    sentinel.inc_ref();
+  }
+
+  return sentinel;
+}
+
+std::pair<std::vector<array>, py::object> tree_flatten_with_structure(
+    py::object tree,
+    bool strict = true) {
+  auto sentinel = structure_sentinel();
+  std::vector<array> flat_tree;
+  auto structure = tree_map(
+      tree,
+      [&flat_tree, sentinel = std::move(sentinel), strict](py::handle obj) {
+        if (py::isinstance<array>(obj)) {
+          flat_tree.push_back(py::cast<array>(obj));
+          return sentinel;
+        } else if (!strict) {
+          return py::cast<py::object>(obj);
+        } else {
+          throw std::invalid_argument(
+              "[tree_flatten] The argument should contain only arrays");
+        }
+      });
+
+  return {flat_tree, structure};
+}
+
+py::object tree_unflatten_from_structure(
+    py::object structure,
+    const std::vector<array>& values,
+    int index = 0) {
+  auto sentinel = structure_sentinel();
+  return tree_map(structure, [&](py::handle obj) {
+    if (obj.is(sentinel)) {
+      return py::cast(values[index++]);
+    } else {
+      return py::cast<py::object>(obj);
+    }
+  });
+}
+
 auto validate_argnums_argnames(
     const std::optional<IntOrVec>& argnums,
     const StrOrVec& argnames) {
@@ -437,6 +485,114 @@ auto py_vmap(
   };
 }
 
+std::unordered_map<size_t, py::object>& tree_cache() {
+  // This map is used to Cache the tree structure of the outputs
+  static std::unordered_map<size_t, py::object> tree_cache_;
+  return tree_cache_;
+}
+
+struct PyCompiledFun {
+  py::function fun;
+  size_t fun_id;
+
+  PyCompiledFun(const py::function& fun)
+      : fun(fun), fun_id(reinterpret_cast<size_t>(fun.ptr())) {}
+
+  PyCompiledFun(const PyCompiledFun&) = delete;
+  PyCompiledFun& operator=(const PyCompiledFun&) = delete;
+  PyCompiledFun& operator=(PyCompiledFun&& other) = delete;
+  PyCompiledFun(PyCompiledFun&& other)
+      : fun(std::move(other.fun)), fun_id(reinterpret_cast<size_t>(fun.ptr())) {
+    other.fun_id = 0;
+  };
+
+  py::object operator()(const py::args& args) {
+    auto compile_fun = [this, &args](const std::vector<array>& a) {
+      // Call the python function and flatten the outputs
+      auto [outputs, py_outputs] = tree_flatten_with_structure(
+          std::move(this->fun(*tree_unflatten(args, a))), true);
+
+      tree_cache().insert({this->fun_id, py_outputs});
+      return outputs;
+    };
+
+    // Inputs must be array or tree of arrays
+    auto inputs = tree_flatten(args, true);
+
+    // Compile and call
+    auto outputs = detail::compile(compile_fun, fun_id)(inputs);
+
+    // Put the outputs back in the container
+    py::object py_outputs = tree_cache().at(fun_id);
+    return tree_unflatten_from_structure(py_outputs, outputs);
+  };
+
+  ~PyCompiledFun() {
+    py::gil_scoped_acquire gil;
+
+    tree_cache().erase(fun_id);
+    detail::compile_erase(fun_id);
+    fun.release().dec_ref();
+  }
+};
+
+class PyCheckpointedFun {
+ public:
+  PyCheckpointedFun(py::function fun) : fun_(std::move(fun)) {}
+
+  ~PyCheckpointedFun() {
+    py::gil_scoped_acquire gil;
+
+    fun_.release().dec_ref();
+  }
+
+  struct InnerFunction {
+    py::object fun_;
+    py::object args_structure_;
+    std::weak_ptr<py::object> output_structure_;
+
+    InnerFunction(
+        py::object fun,
+        py::object args_structure,
+        std::weak_ptr<py::object> output_structure)
+        : fun_(std::move(fun)),
+          args_structure_(std::move(args_structure)),
+          output_structure_(output_structure) {}
+    ~InnerFunction() {
+      py::gil_scoped_acquire gil;
+
+      fun_.release().dec_ref();
+      args_structure_.release().dec_ref();
+    }
+
+    std::vector<array> operator()(const std::vector<array>& inputs) {
+      auto args = py::cast<py::tuple>(
+          tree_unflatten_from_structure(args_structure_, inputs));
+      auto [outputs, output_structure] =
+          tree_flatten_with_structure(fun_(*args[0], **args[1]), false);
+      if (auto s = output_structure_.lock()) {
+        *s = output_structure;
+      }
+      return outputs;
+    }
+  };
+
+  py::object operator()(const py::args& args, const py::kwargs& kwargs) {
+    auto output_structure = std::make_shared<py::object>();
+    auto full_args = py::make_tuple(args, kwargs);
+    auto [inputs, args_structure] =
+        tree_flatten_with_structure(full_args, false);
+
+    auto outputs = checkpoint(
+        InnerFunction(fun_, args_structure, output_structure))(inputs);
+
+    return tree_unflatten_from_structure(*output_structure, outputs);
+  }
+
+ private:
+  py::function fun_;
+};
+
 void init_transforms(py::module_& m) {
   py::options options;
   options.disable_function_signatures();
@@ -445,7 +601,10 @@ void init_transforms(py::module_& m) {
       "eval",
       [](const py::args& args) {
         std::vector<array> arrays = tree_flatten(args);
-        eval(arrays);
+        {
+          py::gil_scoped_release nogil;
+          eval(arrays);
+        }
       },
       R"pbdoc(
         eval(*args) -> None
@@ -679,45 +838,6 @@ void init_transforms(py::module_& m) {
         Returns:
             function: The vectorized function.
       )pbdoc");
-  m.def(
-      "simplify",
-      [](const py::args& args) {
-        std::vector<array> arrays = tree_flatten(args);
-        simplify(arrays);
-      },
-      R"pbdoc(
-        simplify(*args) -> None
-
-        Simplify the graph that computes the arrays.
-
-        Run a few fast graph simplification operations to reuse computation and
-        reduce memory consumption. This function is meant to be run every time
-        so its overhead should be small, approximately 1ms for a graph with a
-        few thousand nodes.
-
-        .. code-block:: python
-
-          import mlx.core as mx
-
-          def foo(x):
-            y = x @ x
-            z = x @ x
-            return y + z
-
-          x = mx.ones((10, 10))
-          y = foo(x)
-          z = foo(x)
-
-          # Computes the matmul twice
-          mx.eval(y)
-
-          # Computes the matmul once
-          mx.simplify(z)
-          mx.eval(z)
-
-        Args:
-          args: Any number of arrays and/or trees of arrays to be simplified.
-      )pbdoc");
   m.def(
       "export_to_dot",
       [](py::object file, const py::args& args) {
@@ -736,4 +856,50 @@ void init_transforms(py::module_& m) {
         }
       },
       "file"_a);
+  m.def(
+      "compile",
+      [](const py::function& fun) {
+        return py::cpp_function(PyCompiledFun{fun});
+      },
+      "fun"_a,
+      R"pbdoc(
+        compile(fun: function) -> function
+
+        Returns a compiled function which produces the same output as ``fun``.
+
+        Args:
+            fun (function): A function which takes a variable number of
+              :class:`array` or trees of :class:`array` and returns
+              a variable number of :class:`array` or trees of :class:`array`.
+
+        Returns:
+            function: A compiled function which has the same input arguments
+            as ``fun`` and returns the the same output(s).
+      )pbdoc");
+  m.def(
+      "disable_compile",
+      &disable_compile,
+      R"pbdoc(
+        disable_compile() -> None
+
+        Globally disable compilation. Setting the environment variable
+        ``MLX_DISABLE_COMPILE`` can also be used to disable compilation.
+      )pbdoc");
+  m.def(
+      "enable_compile",
+      &enable_compile,
+      R"pbdoc(
+        enable_compiler() -> None
+
+        Globally enable compilation. This will override the environment
+        variable ``MLX_DISABLE_COMPILE`` if set.
+      )pbdoc");
+  m.def(
+      "checkpoint",
+      [](py::function fun) { return py::cpp_function(PyCheckpointedFun{fun}); },
+      "fun"_a);
+
+  // Register static Python object cleanup before the interpreter exits
+  auto atexit = py::module_::import("atexit");
+  atexit.attr("register")(py::cpp_function([]() { tree_cache().clear(); }));
 }

python/tests/mlx_tests.py

@@ -65,11 +65,9 @@ class MLXTestCase(unittest.TestCase):
         )
         if not isinstance(mx_res, mx.array) and not isinstance(expected, mx.array):
             np.testing.assert_allclose(mx_res, expected, rtol=rtol, atol=atol)
+            return
         elif not isinstance(mx_res, mx.array):
             mx_res = mx.array(mx_res)
-            self.assertTrue(mx.allclose(mx_res, expected, rtol=rtol, atol=atol))
         elif not isinstance(expected, mx.array):
             expected = mx.array(expected)
-            self.assertTrue(mx.allclose(mx_res, expected, rtol=rtol, atol=atol))
-        else:
-            self.assertTrue(mx.allclose(mx_res, expected, rtol=rtol, atol=atol))
+        self.assertTrue(mx.allclose(mx_res, expected, rtol=rtol, atol=atol))

python/tests/test_array.py

@@ -94,7 +94,7 @@ class TestArray(mlx_tests.MLXTestCase):
         self.assertEqual(x.ndim, 0)
         self.assertEqual(x.itemsize, 4)
         self.assertEqual(x.nbytes, 4)
-        self.assertEqual(x.shape, [])
+        self.assertEqual(x.shape, ())
         self.assertEqual(x.dtype, mx.int32)
         self.assertEqual(x.item(), 1)
         self.assertTrue(isinstance(x.item(), int))
@@ -116,7 +116,7 @@ class TestArray(mlx_tests.MLXTestCase):
         x = mx.array(1.0)
         self.assertEqual(x.size, 1)
         self.assertEqual(x.ndim, 0)
-        self.assertEqual(x.shape, [])
+        self.assertEqual(x.shape, ())
         self.assertEqual(x.dtype, mx.float32)
         self.assertEqual(x.item(), 1.0)
         self.assertTrue(isinstance(x.item(), float))
@@ -124,14 +124,14 @@ class TestArray(mlx_tests.MLXTestCase):
         x = mx.array(False)
         self.assertEqual(x.size, 1)
         self.assertEqual(x.ndim, 0)
-        self.assertEqual(x.shape, [])
+        self.assertEqual(x.shape, ())
         self.assertEqual(x.dtype, mx.bool_)
         self.assertEqual(x.item(), False)
         self.assertTrue(isinstance(x.item(), bool))
 
         x = mx.array(complex(1, 1))
         self.assertEqual(x.ndim, 0)
-        self.assertEqual(x.shape, [])
+        self.assertEqual(x.shape, ())
         self.assertEqual(x.dtype, mx.complex64)
         self.assertEqual(x.item(), complex(1, 1))
         self.assertTrue(isinstance(x.item(), complex))
@@ -139,7 +139,7 @@ class TestArray(mlx_tests.MLXTestCase):
         x = mx.array([True, False, True])
         self.assertEqual(x.dtype, mx.bool_)
         self.assertEqual(x.ndim, 1)
-        self.assertEqual(x.shape, [3])
+        self.assertEqual(x.shape, (3,))
         self.assertEqual(len(x), 3)
 
         x = mx.array([True, False, True], mx.float32)
@@ -148,7 +148,7 @@ class TestArray(mlx_tests.MLXTestCase):
         x = mx.array([0, 1, 2])
         self.assertEqual(x.dtype, mx.int32)
         self.assertEqual(x.ndim, 1)
-        self.assertEqual(x.shape, [3])
+        self.assertEqual(x.shape, (3,))
 
         x = mx.array([0, 1, 2], mx.float32)
         self.assertEqual(x.dtype, mx.float32)
@@ -156,12 +156,12 @@ class TestArray(mlx_tests.MLXTestCase):
         x = mx.array([0.0, 1.0, 2.0])
         self.assertEqual(x.dtype, mx.float32)
         self.assertEqual(x.ndim, 1)
-        self.assertEqual(x.shape, [3])
+        self.assertEqual(x.shape, (3,))
 
         x = mx.array([1j, 1 + 0j])
         self.assertEqual(x.dtype, mx.complex64)
         self.assertEqual(x.ndim, 1)
-        self.assertEqual(x.shape, [2])
+        self.assertEqual(x.shape, (2,))
 
         # From tuple
         x = mx.array((1, 2, 3), mx.int32)
@@ -181,17 +181,17 @@ class TestArray(mlx_tests.MLXTestCase):
     def test_construction_from_lists(self):
         x = mx.array([])
         self.assertEqual(x.size, 0)
-        self.assertEqual(x.shape, [0])
+        self.assertEqual(x.shape, (0,))
         self.assertEqual(x.dtype, mx.float32)
 
         x = mx.array([[], [], []])
         self.assertEqual(x.size, 0)
-        self.assertEqual(x.shape, [3, 0])
+        self.assertEqual(x.shape, (3, 0))
         self.assertEqual(x.dtype, mx.float32)
 
         x = mx.array([[[], []], [[], []], [[], []]])
         self.assertEqual(x.size, 0)
-        self.assertEqual(x.shape, [3, 2, 0])
+        self.assertEqual(x.shape, (3, 2, 0))
         self.assertEqual(x.dtype, mx.float32)
 
         # Check failure cases
@@ -226,6 +226,14 @@ class TestArray(mlx_tests.MLXTestCase):
         x = mx.array([1 + 0j, 2j, True, 0], mx.complex64)
         self.assertEqual(x.tolist(), [1 + 0j, 2j, 1 + 0j, 0j])
 
+        xnp = np.array([0, 4294967295], dtype=np.uint32)
+        x = mx.array([0, 4294967295], dtype=mx.uint32)
+        self.assertTrue(np.array_equal(x, xnp))
+
+        xnp = np.array([0, 4294967295], dtype=np.float32)
+        x = mx.array([0, 4294967295], dtype=mx.float32)
+        self.assertTrue(np.array_equal(x, xnp))
+
     def test_construction_from_lists_of_mlx_arrays(self):
         dtypes = [
             mx.bool_,
@@ -428,19 +436,19 @@ class TestArray(mlx_tests.MLXTestCase):
         a = np.array([])
         x = mx.array(a)
         self.assertEqual(x.size, 0)
-        self.assertEqual(x.shape, [0])
+        self.assertEqual(x.shape, (0,))
         self.assertEqual(x.dtype, mx.float32)
 
         a = np.array([[], [], []])
         x = mx.array(a)
         self.assertEqual(x.size, 0)
-        self.assertEqual(x.shape, [3, 0])
+        self.assertEqual(x.shape, (3, 0))
         self.assertEqual(x.dtype, mx.float32)
 
         a = np.array([[[], []], [[], []], [[], []]])
         x = mx.array(a)
         self.assertEqual(x.size, 0)
-        self.assertEqual(x.shape, [3, 2, 0])
+        self.assertEqual(x.shape, (3, 2, 0))
         self.assertEqual(x.dtype, mx.float32)
 
         # Content test
@@ -448,7 +456,7 @@ class TestArray(mlx_tests.MLXTestCase):
         x = mx.array(a)
         self.assertEqual(x.dtype, mx.float32)
         self.assertEqual(x.ndim, 3)
-        self.assertEqual(x.shape, [3, 5, 4])
+        self.assertEqual(x.shape, (3, 5, 4))
 
         y = np.asarray(x)
         self.assertTrue(np.allclose(a, y))
@@ -457,7 +465,7 @@ class TestArray(mlx_tests.MLXTestCase):
         x = mx.array(a)
         self.assertEqual(x.dtype, mx.int32)
         self.assertEqual(x.ndim, 0)
-        self.assertEqual(x.shape, [])
+        self.assertEqual(x.shape, ())
         self.assertEqual(x.item(), 3)
 
         # mlx to numpy test
@@ -475,7 +483,7 @@ class TestArray(mlx_tests.MLXTestCase):
         x = np.array(cvals)
         y = mx.array(x)
         self.assertEqual(y.dtype, mx.complex64)
-        self.assertEqual(y.shape, [3])
+        self.assertEqual(y.shape, (3,))
         self.assertEqual(y.tolist(), cvals)
 
         y = mx.array([0j, 1, 1 + 1j])

python/tests/test_autograd.py

@@ -380,6 +380,19 @@ class TestAutograd(mlx_tests.MLXTestCase):
             out = mx.grad(fun)(mx.array(1.0, t), mx.array(1.0, t))
             self.assertEqual(out.dtype, t)
 
+    def test_power_grad(self):
+        x = mx.array(0.0)
+        g = mx.grad(lambda x: x**2)(x)
+        self.assertEqual(g.item(), 0.0)
+
+        x = mx.array(0.0)
+        g = mx.grad(lambda x: x**1.5)(x)
+        self.assertEqual(g.item(), 0.0)
+
+        x = mx.array(2.0)
+        g = mx.grad(lambda x: x**2)(x)
+        self.assertAlmostEqual(g.item(), 4.0)
+
 
 if __name__ == "__main__":
     unittest.main()

python/tests/test_blas.py

@@ -576,8 +576,12 @@ class TestBlas(mlx_tests.MLXTestCase):
                 ],
             )
 
-            self.assertTrue(mx.allclose(out_ref[0], out_test[0], atol=1e-5).item())
+            self.assertTrue(mx.allclose(out_ref[0], out_test[0], atol=1e-4).item())
 
             for r, t in zip(dout_ref, dout_test):
-                self.assertListEqual(r.shape, t.shape)
-                self.assertTrue(mx.allclose(r, t, atol=1e-5).item())
+                self.assertEqual(r.shape, t.shape)
+                self.assertTrue(mx.allclose(r, t, atol=1e-4).item())
+
+
+if __name__ == "__main__":
+    unittest.main()

python/tests/test_compile.py

@@ -0,0 +1,195 @@
+# Copyright © 2023-2024 Apple Inc.
+
+import io
+import unittest
+
+import mlx.core as mx
+import mlx_tests
+
+
+class TestCompile(mlx_tests.MLXTestCase):
+    def test_simple_compile(self):
+        def fun(x, y):
+            return x + y
+
+        compiled_fn = mx.compile(fun)
+        compiled_fn = mx.compile(fun)
+        x = mx.array(1.0)
+        y = mx.array(1.0)
+        out = compiled_fn(x, y)
+        self.assertEqual(out.item(), 2.0)
+
+        # Try again
+        out = compiled_fn(x, y)
+        self.assertEqual(out.item(), 2.0)
+
+        # Change sizes
+        x = mx.array([1.0, 2.0])
+        out = compiled_fn(x, y)
+        self.assertTrue(mx.array_equal(out, mx.array([2.0, 3.0])))
+
+        y = mx.array([1.0, 2.0])
+        out = compiled_fn(x, y)
+        self.assertTrue(mx.array_equal(out, mx.array([2.0, 4.0])))
+
+        # Change types
+        x = mx.array([1, 2], mx.int32)
+        y = mx.array([1, 2], mx.int32)
+        out = compiled_fn(x, y)
+        self.assertEqual(out.dtype, mx.int32)
+        self.assertTrue(mx.array_equal(out, mx.array([2, 4])))
+
+    def test_compile_grad(self):
+        def loss_fn(x):
+            return mx.exp(x).sum()
+
+        grad_fn = mx.grad(loss_fn)
+
+        x = mx.array([0.5, -0.5, 1.2])
+        dfdx = grad_fn(x)
+        compile_grad_fn = mx.compile(grad_fn)
+        c_dfdx = grad_fn(x)
+
+        self.assertTrue(mx.allclose(c_dfdx, dfdx))
+
+        # Run it again without calling compile
+        c_dfdx = compile_grad_fn(x)
+        self.assertTrue(mx.allclose(c_dfdx, dfdx))
+
+        # Run it again with calling compile
+        c_dfdx = mx.compile(grad_fn)(x)
+        self.assertTrue(mx.allclose(c_dfdx, dfdx))
+
+        # Value and grad
+        def loss_fn(x):
+            return mx.exp(x).sum(), mx.sin(x)
+
+        val_and_grad_fn = mx.value_and_grad(loss_fn)
+        (loss, val), dfdx = val_and_grad_fn(x)
+        (c_loss, c_val), c_dfdx = mx.compile(val_and_grad_fn)(x)
+
+        self.assertTrue(mx.allclose(c_dfdx, dfdx))
+        self.assertTrue(mx.allclose(c_loss, loss))
+        self.assertTrue(mx.allclose(c_val, val))
+
+    def test_compile_inputs_with_primitives(self):
+        x = mx.array([1, 2, 3])
+        y = mx.array([1, 2, 3])
+        for _ in range(5):
+            x = x + y
+            y = y + 1
+
+        def fun(x, y):
+            return x * y
+
+        out = fun(x, y)
+
+        x = mx.array([1, 2, 3])
+        y = mx.array([1, 2, 3])
+        for _ in range(5):
+            x = x + y
+            y = y + 1
+
+        c_out = mx.compile(fun)(x, y)
+        self.assertTrue(mx.array_equal(out, c_out))
+
+        # Try again
+        c_out = mx.compile(fun)(x, y)
+        self.assertTrue(mx.array_equal(out, c_out))
+
+    def test_compile_with_closure(self):
+        x = mx.array(1)
+
+        def closure(y):
+            return x + y
+
+        compiled = mx.compile(closure)
+        out = compiled(mx.array(1))
+        self.assertEqual(out.item(), 2)
+
+        # Try again
+        out = compiled(mx.array(1))
+        self.assertEqual(out.item(), 2)
+
+        # Change the shape of the enclosed variable
+        x = mx.array([1, 2])
+        out = compiled(mx.array(1))
+
+        # We still get the original input (closures are not updated)
+        self.assertEqual(out.item(), 2)
+
+        # Try with a tree of enclosed variables
+        x = {"a": mx.array(1), "b": mx.array(2)}
+
+        def closure(y):
+            return x["a"] + y + x["b"]
+
+        compiled = mx.compile(closure)
+        out = compiled(mx.array(1))
+        self.assertEqual(out.item(), 4)
+
+        # Change the shape of one input
+        x["a"] = mx.array([4, 5])
+        out = compiled(mx.array(1))
+        self.assertEqual(out.item(), 4)
+
+        x["b"] = mx.array([-6, -8])
+        out = compiled(mx.array(1))
+        self.assertEqual(out.item(), 4)
+
+        # Enclosed variable is not evaluated yet
+        x = mx.array(1)
+        x = x + x
+
+        def closure(y):
+            return x + y
+
+        compiled = mx.compile(closure)
+        out = compiled(mx.array(2))
+        self.assertEqual(out.item(), 4)
+
+        # And again
+        out = compiled(mx.array(2))
+        self.assertEqual(out.item(), 4)
+
+    def test_function_creates_array(self):
+        def fun(x):
+            return x + mx.array(1)
+
+        cfun = mx.compile(fun)
+        out = cfun(mx.array(3))
+        self.assertEqual(out.item(), 4)
+
+        # And again
+        out = cfun(mx.array(3))
+        self.assertEqual(out.item(), 4)
+
+    def test_enable_disable(self):
+        def fun(x):
+            y = x + 1
+            z = x + 1
+            return y + z
+
+        def count_prims(outputs):
+            buf = io.StringIO()
+            mx.export_to_dot(buf, outputs)
+            buf.seek(0)
+            return len([l for l in buf.read().split() if "label" in l])
+
+        x = mx.array(1.0)
+        cfun = mx.compile(fun)
+        n_compiled = count_prims(cfun(x))
+
+        # Check disabled
+        mx.disable_compile()
+        n_uncompiled = count_prims(cfun(x))
+        self.assertTrue(n_compiled < n_uncompiled)
+
+        # Check renabled
+        mx.enable_compile()
+        n_enable_compiled = count_prims(cfun(x))
+        self.assertEqual(n_compiled, n_enable_compiled)
+
+
+if __name__ == "__main__":
+    unittest.main()

python/tests/test_conv.py

@@ -44,7 +44,7 @@ class TestConv(mlx_tests.MLXTestCase):
                     c_np = np.convolve(a_np, v_np, mode=mode)
                     c_mx = mx.convolve(a_mx, v_mx, mode=mode)
 
-                    self.assertListEqual(list(c_mx.shape), list(c_np.shape))
+                    self.assertEqual(c_mx.shape, c_np.shape)
                     self.assertTrue(np.allclose(c_mx, c_np, atol=atol))
 
     @unittest.skipIf(not has_torch, "requires Torch")
@@ -102,7 +102,7 @@ class TestConv(mlx_tests.MLXTestCase):
                 )
                 out_pt = torch.transpose(out_pt, 2, 1)
 
-                self.assertListEqual(list(out_pt.shape), out_mx.shape)
+                self.assertEqual(out_pt.shape, out_mx.shape)
                 self.assertTrue(np.allclose(out_pt.numpy(), out_mx, atol=atol))
 
         for dtype in ("float32",):
@@ -141,7 +141,7 @@ class TestConv(mlx_tests.MLXTestCase):
                 out_pt = torch.conv1d(in_pt, wt_pt)
                 out_pt = torch.transpose(out_pt, 2, 1)
 
-                self.assertListEqual(list(out_pt.shape), out_mx.shape)
+                self.assertEqual(out_pt.shape, out_mx.shape)
                 self.assertTrue(np.allclose(out_pt.numpy(), out_mx, atol=1e-5))
 
     @unittest.skipIf(not has_torch, "requires Torch")
@@ -228,12 +228,12 @@ class TestConv(mlx_tests.MLXTestCase):
 
                 mx_grad_in, mx_grad_wt = outs_mx
 
-                self.assertListEqual(list(pt_grad_in.shape), mx_grad_in.shape)
-                self.assertListEqual(list(in_mx.shape), mx_grad_in.shape)
+                self.assertEqual(pt_grad_in.shape, mx_grad_in.shape)
+                self.assertEqual(in_mx.shape, mx_grad_in.shape)
                 self.assertTrue(np.allclose(pt_grad_in, mx_grad_in, atol=atol))
 
-                self.assertListEqual(list(pt_grad_wt.shape), mx_grad_wt.shape)
-                self.assertListEqual(list(wt_mx.shape), mx_grad_wt.shape)
+                self.assertEqual(pt_grad_wt.shape, mx_grad_wt.shape)
+                self.assertEqual(wt_mx.shape, mx_grad_wt.shape)
                 self.assertTrue(np.allclose(pt_grad_wt, mx_grad_wt, atol=atol))
 
         for dtype in ("float32",):
@@ -309,7 +309,7 @@ class TestConv(mlx_tests.MLXTestCase):
                 )
                 out_pt = torch.permute(out_pt, (0, 2, 3, 1)).numpy(force=True)
 
-                self.assertListEqual(list(out_pt.shape), list(out_mx.shape))
+                self.assertEqual(out_pt.shape, out_mx.shape)
                 self.assertTrue(np.allclose(out_pt, out_mx, atol=atol))
 
         for dtype in ("float32",):
@@ -419,12 +419,12 @@ class TestConv(mlx_tests.MLXTestCase):
 
                 mx_grad_in, mx_grad_wt = outs_mx
 
-                self.assertListEqual(list(pt_grad_in.shape), mx_grad_in.shape)
-                self.assertListEqual(list(in_mx.shape), mx_grad_in.shape)
+                self.assertEqual(pt_grad_in.shape, mx_grad_in.shape)
+                self.assertEqual(in_mx.shape, mx_grad_in.shape)
                 self.assertTrue(np.allclose(pt_grad_in, mx_grad_in, atol=atol))
 
-                self.assertListEqual(list(pt_grad_wt.shape), mx_grad_wt.shape)
-                self.assertListEqual(list(wt_mx.shape), mx_grad_wt.shape)
+                self.assertEqual(pt_grad_wt.shape, mx_grad_wt.shape)
+                self.assertEqual(wt_mx.shape, mx_grad_wt.shape)
                 self.assertTrue(np.allclose(pt_grad_wt, mx_grad_wt, atol=atol))
 
         for dtype in ("float32",):

python/tests/test_init.py

@@ -0,0 +1,94 @@
+# Copyright © 2023 Apple Inc.
+import unittest
+
+import mlx.core as mx
+import mlx.nn.init as init
+import mlx_tests
+import numpy as np
+
+
+class TestInit(mlx_tests.MLXTestCase):
+    def test_constant(self):
+        value = 5.0
+
+        for dtype in [mx.float32, mx.float16]:
+            initializer = init.constant(value, dtype)
+            for shape in [(3,), (3, 3), (3, 3, 3)]:
+                result = initializer(mx.array(mx.zeros(shape)))
+                with self.subTest(shape=shape):
+                    self.assertEqual(result.shape, shape)
+                    self.assertEqual(result.dtype, dtype)
+
+    def test_normal(self):
+        mean = 0.0
+        std = 1.0
+        for dtype in [mx.float32, mx.float16]:
+            initializer = init.normal(mean, std, dtype=dtype)
+            for shape in [(3,), (3, 3), (3, 3, 3)]:
+                result = initializer(mx.array(np.empty(shape)))
+                with self.subTest(shape=shape):
+                    self.assertEqual(result.shape, shape)
+                    self.assertEqual(result.dtype, dtype)
+
+    def test_uniform(self):
+        low = -1.0
+        high = 1.0
+
+        for dtype in [mx.float32, mx.float16]:
+            initializer = init.uniform(low, high, dtype)
+            for shape in [(3,), (3, 3), (3, 3, 3)]:
+                result = initializer(mx.array(np.empty(shape)))
+                with self.subTest(shape=shape):
+                    self.assertEqual(result.shape, shape)
+                    self.assertEqual(result.dtype, dtype)
+                    self.assertTrue(mx.all(result >= low) and mx.all(result <= high))
+
+    def test_identity(self):
+        for dtype in [mx.float32, mx.float16]:
+            initializer = init.identity(dtype)
+            for shape in [(3,), (3, 3), (3, 3, 3)]:
+                result = initializer(mx.zeros((3, 3)))
+                self.assertTrue(mx.array_equal(result, mx.eye(3)))
+                self.assertEqual(result.dtype, dtype)
+                with self.assertRaises(ValueError):
+                    result = initializer(mx.zeros((3, 2)))
+
+    def test_glorot_normal(self):
+        for dtype in [mx.float32, mx.float16]:
+            initializer = init.glorot_normal(dtype)
+            for shape in [(3, 3), (3, 3, 3)]:
+                result = initializer(mx.array(np.empty(shape)))
+                with self.subTest(shape=shape):
+                    self.assertEqual(result.shape, shape)
+                    self.assertEqual(result.dtype, dtype)
+
+    def test_glorot_uniform(self):
+        for dtype in [mx.float32, mx.float16]:
+            initializer = init.glorot_uniform(dtype)
+            for shape in [(3, 3), (3, 3, 3)]:
+                result = initializer(mx.array(np.empty(shape)))
+                with self.subTest(shape=shape):
+                    self.assertEqual(result.shape, shape)
+                    self.assertEqual(result.dtype, dtype)
+
+    def test_he_normal(self):
+        for dtype in [mx.float32, mx.float16]:
+            initializer = init.he_normal(dtype)
+            for shape in [(3, 3), (3, 3, 3)]:
+                result = initializer(mx.array(np.empty(shape)))
+                with self.subTest(shape=shape):
+                    self.assertEqual(result.shape, shape)
+                    self.assertEqual(result.dtype, dtype)
+
+    def test_he_uniform(self):
+        for dtype in [mx.float32, mx.float16]:
+            initializer = init.he_uniform(dtype)
+            for shape in [(3, 3), (3, 3, 3)]:
+                result = initializer(mx.array(np.empty(shape)))
+                with self.subTest(shape=shape):
+                    self.assertEqual(result.shape, shape)
+                    self.assertEqual(result.dtype, dtype)
+
+
+if __name__ == "__main__":
+    unittest.main()

python/tests/test_linalg.py

@@ -89,6 +89,37 @@ class TestLinalg(mlx_tests.MLXTestCase):
         out_mx = mx.linalg.norm(x_mx, ord="fro")
         self.assertTrue(np.allclose(out_np, out_mx, atol=1e-5, rtol=1e-6))
 
+    def test_qr_factorization(self):
+        with self.assertRaises(ValueError):
+            mx.linalg.qr(mx.array(0.0))
+
+        with self.assertRaises(ValueError):
+            mx.linalg.qr(mx.array([0.0, 1.0]))
+
+        with self.assertRaises(ValueError):
+            mx.linalg.qr(mx.array([[0, 1], [1, 0]]))
+
+        A = mx.array([[2.0, 3.0], [1.0, 2.0]])
+        Q, R = mx.linalg.qr(A, stream=mx.cpu)
+        out = Q @ R
+        self.assertTrue(mx.allclose(out, A))
+        out = Q @ Q
+        self.assertTrue(mx.allclose(out, mx.eye(2), rtol=1e-5, atol=1e-7))
+        self.assertTrue(mx.allclose(mx.tril(R, -1), mx.zeros_like(R)))
+        self.assertEqual(Q.dtype, mx.float32)
+        self.assertEqual(R.dtype, mx.float32)
+
+        # Multiple matrices
+        B = mx.array([[-1.0, 2.0], [-4.0, 1.0]])
+        A = mx.stack([A, B])
+        Q, R = mx.linalg.qr(A, stream=mx.cpu)
+        for a, q, r in zip(A, Q, R):
+            out = q @ r
+            self.assertTrue(mx.allclose(out, a))
+            out = q @ q
+            self.assertTrue(mx.allclose(out, mx.eye(2), rtol=1e-5, atol=1e-7))
+            self.assertTrue(mx.allclose(mx.tril(r, -1), mx.zeros_like(r)))
+
 
 if __name__ == "__main__":
     unittest.main()

python/tests/test_load.py

@@ -117,6 +117,115 @@ class TestLoad(mlx_tests.MLXTestCase):
                             mx.array_equal(load_dict["test"], save_dict["test"])
                         )
 
+    def test_save_and_load_gguf_metadata_basic(self):
+        if not os.path.isdir(self.test_dir):
+            os.mkdir(self.test_dir)
+
+        save_file_mlx = os.path.join(self.test_dir, f"mlx_gguf_with_metadata.gguf")
+        save_dict = {"test": mx.ones((4, 4), dtype=mx.int32)}
+        metadata = {}
+
+        # Empty works
+        mx.save_gguf(save_file_mlx, save_dict, metadata)
+
+        # Loads without the metadata
+        load_dict = mx.load(save_file_mlx)
+        self.assertTrue("test" in load_dict)
+        self.assertTrue(mx.array_equal(load_dict["test"], save_dict["test"]))
+
+        # Loads empty metadata
+        load_dict, meta_load_dict = mx.load(save_file_mlx, return_metadata=True)
+        self.assertTrue("test" in load_dict)
+        self.assertTrue(mx.array_equal(load_dict["test"], save_dict["test"]))
+        self.assertEqual(len(meta_load_dict), 0)
+
+        # Loads string metadata
+        metadata = {"meta": "data"}
+        mx.save_gguf(save_file_mlx, save_dict, metadata)
+        load_dict, meta_load_dict = mx.load(save_file_mlx, return_metadata=True)
+        self.assertTrue("test" in load_dict)
+        self.assertTrue(mx.array_equal(load_dict["test"], save_dict["test"]))
+        self.assertEqual(len(meta_load_dict), 1)
+        self.assertTrue("meta" in meta_load_dict)
+        self.assertEqual(meta_load_dict["meta"], "data")
+
+    def test_save_and_load_gguf_metadata_arrays(self):
+        if not os.path.isdir(self.test_dir):
+            os.mkdir(self.test_dir)
+
+        save_file_mlx = os.path.join(self.test_dir, f"mlx_gguf_with_metadata.gguf")
+        save_dict = {"test": mx.ones((4, 4), dtype=mx.int32)}
+
+        # Test scalars and one dimensional arrays
+        for t in [
+            mx.uint8,
+            mx.int8,
+            mx.uint16,
+            mx.int16,
+            mx.uint32,
+            mx.int32,
+            mx.uint64,
+            mx.int64,
+            mx.float32,
+        ]:
+            for shape in [(), (2,)]:
+                arr = mx.random.uniform(shape=shape).astype(t)
+                metadata = {"meta": arr}
+                mx.save_gguf(save_file_mlx, save_dict, metadata)
+                _, meta_load_dict = mx.load(save_file_mlx, return_metadata=True)
+                self.assertEqual(len(meta_load_dict), 1)
+                self.assertTrue("meta" in meta_load_dict)
+                self.assertTrue(mx.array_equal(meta_load_dict["meta"], arr))
+                self.assertEqual(meta_load_dict["meta"].dtype, arr.dtype)
+
+        for t in [mx.float16, mx.bfloat16, mx.complex64]:
+            with self.assertRaises(ValueError):
+                arr = mx.array(1, t)
+                metadata = {"meta": arr}
+                mx.save_gguf(save_file_mlx, save_dict, metadata)
+
+    def test_save_and_load_gguf_metadata_mixed(self):
+        if not os.path.isdir(self.test_dir):
+            os.mkdir(self.test_dir)
+
+        save_file_mlx = os.path.join(self.test_dir, f"mlx_gguf_with_metadata.gguf")
+        save_dict = {"test": mx.ones((4, 4), dtype=mx.int32)}
+
+        # Test string and array
+        arr = mx.array(1.5)
+        metadata = {"meta1": arr, "meta2": "data"}
+        mx.save_gguf(save_file_mlx, save_dict, metadata)
+        _, meta_load_dict = mx.load(save_file_mlx, return_metadata=True)
+        self.assertEqual(len(meta_load_dict), 2)
+        self.assertTrue("meta1" in meta_load_dict)
+        self.assertTrue(mx.array_equal(meta_load_dict["meta1"], arr))
+        self.assertEqual(meta_load_dict["meta1"].dtype, arr.dtype)
+        self.assertTrue("meta2" in meta_load_dict)
+        self.assertEqual(meta_load_dict["meta2"], "data")
+
+        # Test list of strings
+        metadata = {"meta": ["data1", "data2", "data345"]}
+        mx.save_gguf(save_file_mlx, save_dict, metadata)
+        _, meta_load_dict = mx.load(save_file_mlx, return_metadata=True)
+        self.assertEqual(len(meta_load_dict), 1)
+        self.assertEqual(meta_load_dict["meta"], metadata["meta"])
+
+        # Test a combination of stuff
+        metadata = {
+            "meta1": ["data1", "data2", "data345"],
+            "meta2": mx.array([1, 2, 3, 4]),
+            "meta3": "data",
+            "meta4": mx.array(1.5),
+        }
+        mx.save_gguf(save_file_mlx, save_dict, metadata)
+        _, meta_load_dict = mx.load(save_file_mlx, return_metadata=True)
+        self.assertEqual(len(meta_load_dict), 4)
+        for k, v in metadata.items():
+            if isinstance(v, mx.array):
+                self.assertTrue(mx.array_equal(meta_load_dict[k], v))
+            else:
+                self.assertEqual(meta_load_dict[k], v)
+
     def test_save_and_load_fs(self):
         if not os.path.isdir(self.test_dir):
             os.mkdir(self.test_dir)

python/tests/test_losses.py

@@ -10,100 +10,115 @@ import numpy as np
 
 class TestLosses(mlx_tests.MLXTestCase):
     def test_cross_entropy(self):
+        # No weights, no label smoothing
         logits = mx.array([[0.0, -float("inf")], [-float("inf"), 0.0]])
-        targets = mx.array([0, 1])
+        indices = mx.array([0, 1])
+        expected = mx.array([0.0, 0.0])
+        loss = nn.losses.cross_entropy(logits, indices, reduction="none")
+        self.assertTrue(mx.allclose(loss, expected))
 
-        # Test with reduction 'none'
-        losses_none = nn.losses.cross_entropy(logits, targets, reduction="none")
-        expected_none = mx.array([0.0, 0.0])
-        self.assertTrue(mx.array_equal(losses_none, expected_none))
+        probs = mx.array([[1.0, 0.0], [0.0, 1.0]])
+        loss = nn.losses.cross_entropy(logits, probs, reduction="none")
+        self.assertTrue(mx.isnan(loss).all())  # produce NaNs, like PyTorch
 
-        # Test with reduction 'mean'
-        losses_mean = nn.losses.cross_entropy(logits, targets, reduction="mean")
-        expected_mean = mx.mean(expected_none)
-        self.assertEqual(losses_mean, expected_mean)
-
-        # Test with reduction 'sum'
-        losses_sum = nn.losses.cross_entropy(logits, targets, reduction="sum")
-        expected_sum = mx.sum(expected_none)
-        self.assertEqual(losses_sum, expected_sum)
-
-        # Test cases with weights and no label smoothing
+        # With weights, no label smoothing
         logits = mx.array([[2.0, -1.0], [-1.0, 2.0]])
-        targets = mx.array([0, 1])
+        indices = mx.array([0, 1])
         weights = mx.array([1.0, 2.0])
+        expected = mx.array([0.04858735, 0.0971747])
+        loss = nn.losses.cross_entropy(
+            logits, indices, weights=weights, reduction="none"
+        )
+        self.assertTrue(mx.allclose(loss, expected))
 
-        # Reduction 'none'
-        losses_none = nn.losses.cross_entropy(
-            logits,
-            targets,
-            weights=weights,
-            reduction="none",
-        )
-        expected_none = mx.array([0.04858735, 0.0971747])  # Calculated losses
-        self.assertTrue(
-            np.allclose(losses_none, expected_none, atol=1e-5),
-            "Test case failed for cross_entropy loss --reduction='none' --weights=[1.0, 2.0]",
-        )
+        probs = mx.array([[1.0, 0.0], [0.0, 1.0]])
+        loss = nn.losses.cross_entropy(logits, probs, weights=weights, reduction="none")
+        self.assertTrue(mx.allclose(loss, expected))
 
-        # Reduction 'mean'
-        losses_mean = nn.losses.cross_entropy(
-            logits,
-            targets,
-            weights=weights,
-            reduction="mean",
-        )
-        expected_mean = mx.mean(expected_none)
-        self.assertTrue(
-            np.allclose(losses_mean, expected_mean, atol=1e-5),
-            "Test case failed for cross_entropy loss --reduction='mean' --weights=[1.0, 2.0]",
+        # No weights, with label smoothing
+        logits = mx.array([[2.0, -1.0], [-1.0, 2.0]])
+        indices = mx.array([0, 1])
+        expected = mx.array([0.498587, 0.498587])
+        loss = nn.losses.cross_entropy(
+            logits, indices, label_smoothing=0.3, reduction="none"
         )
+        self.assertTrue(mx.allclose(loss, expected))
 
-        # Reduction 'sum'
-        losses_sum = nn.losses.cross_entropy(
-            logits,
-            targets,
-            weights=weights,
-            reduction="sum",
-        )
-        expected_sum = mx.sum(expected_none)
-        self.assertTrue(
-            np.allclose(losses_sum, expected_sum, atol=1e-5),
-            "Test case failed for cross_entropy loss --reduction='sum' --weights=[1.0, 2.0]",
+        probs = mx.array([[1.0, 0.0], [0.0, 1.0]])
+        loss = nn.losses.cross_entropy(
+            logits, probs, label_smoothing=0.3, reduction="none"
         )
+        self.assertTrue(mx.allclose(loss, expected))
 
-        # Test case with equal weights and label smoothing > 0
-        logits = mx.array(
-            [[0, 0.2, 0.7, 0.1, 0], [0, 0.9, 0.2, 0.2, 1], [1, 0.2, 0.7, 0.9, 1]]
+        # With weights and label smoothing
+        logits = mx.array([[2.0, -1.0], [-1.0, 2.0]])
+        indices = mx.array([0, 1])
+        weights = mx.array([1.0, 2.0])
+        expected = mx.array([0.49858734, 0.9971747])
+        loss = nn.losses.cross_entropy(
+            logits, indices, weights=weights, label_smoothing=0.3, reduction="none"
         )
-        target = mx.array([2, 1, 0])
+        self.assertTrue(mx.allclose(loss, expected))
 
-        losses_none = nn.losses.cross_entropy(
-            logits, target, label_smoothing=0.3, reduction="none"
-        )
-        expected_none = mx.array([1.29693, 1.38617, 1.48176])
-        self.assertTrue(
-            mx.allclose(expected_none, losses_none),
-            "Test case failed for cross_entropy --label_smoothing=0.3 --reduction='none'",
+        probs = mx.array([[1.0, 0.0], [0.0, 1.0]])
+        loss = nn.losses.cross_entropy(
+            logits, probs, weights=weights, label_smoothing=0.3, reduction="none"
         )
+        self.assertTrue(mx.allclose(loss, expected))
 
-        expected_mean = mx.mean(expected_none)
-        losses_mean = nn.losses.cross_entropy(
-            logits, target, label_smoothing=0.3, reduction="mean"
-        )
-        self.assertTrue(
-            mx.allclose(losses_mean, expected_mean),
-            "Test case failed for cross_entropy --label_smoothing=0.3 --reduction='mean'",
-        )
+    def test_binary_cross_entropy(self):
+        def _test_logits_as_inputs():
+            logits = mx.array([0.105361, 0.223144, 1.20397, 0.916291])
+            targets = mx.array([0, 0, 1, 1])
 
-        expected_sum = mx.sum(expected_none)
-        losses_sum = nn.losses.cross_entropy(
-            logits, target, label_smoothing=0.3, reduction="sum"
-        )
-        self.assertTrue(
-            mx.allclose(losses_sum, expected_sum),
-            "Test case failed for cross_entropy --label_smoothing=0.3 --reduction='sum'",
-        )
+            # Test with reduction 'none'
+            losses_none = nn.losses.binary_cross_entropy(
+                logits, targets, reduction="none"
+            )
+            expected_none = mx.array([0.747215, 0.810930, 0.262365, 0.336472])
+            self.assertTrue(mx.allclose(losses_none, expected_none))
+
+            # Test with reduction 'mean'
+            losses_mean = nn.losses.binary_cross_entropy(
+                logits, targets, reduction="mean"
+            )
+            expected_mean = mx.mean(expected_none)
+            self.assertEqual(losses_mean, expected_mean)
+
+            # Test with reduction 'sum'
+            losses_sum = nn.losses.binary_cross_entropy(
+                logits, targets, reduction="sum"
+            )
+            expected_sum = mx.sum(expected_none)
+            self.assertEqual(losses_sum, expected_sum)
+
+        def _test_probs_as_inputs():
+            probs = mx.array([0.5, 0.6, 0.7, 0.8])
+            targets = mx.array([0, 0, 1, 1])
+
+            # Test with reduction 'none'
+            losses_none = nn.losses.binary_cross_entropy(
+                probs, targets, with_logits=False, reduction="none"
+            )
+            expected_none = mx.array([0.693147, 0.916291, 0.356675, 0.223144])
+            self.assertTrue(mx.allclose(losses_none, expected_none))
+
+            # Test with reduction 'mean'
+            losses_mean = nn.losses.binary_cross_entropy(
+                probs, targets, with_logits=False, reduction="mean"
+            )
+            expected_mean = mx.mean(expected_none)
+            self.assertTrue(mx.allclose(losses_mean, expected_mean))
+
+            # Test with reduction 'sum'
+            losses_sum = nn.losses.binary_cross_entropy(
+                probs, targets, with_logits=False, reduction="sum"
+            )
+            expected_sum = mx.sum(expected_none)
+            self.assertTrue(mx.allclose(losses_sum, expected_sum))
+
+        _test_logits_as_inputs()
+        _test_probs_as_inputs()
 
     def test_l1_loss(self):
         predictions = mx.array([0.5, 0.2, 0.9, 0.0])

python/tests/test_nn.py

@@ -54,16 +54,31 @@ class TestBase(mlx_tests.MLXTestCase):
 
         m.apply_to_modules(assert_training)
 
-    def test_io(self):
+    def test_save_npz_weights(self):
         def make_model():
             return nn.Sequential(nn.Linear(2, 2), nn.ReLU(), nn.Linear(2, 2))
 
         m = make_model()
         tdir = tempfile.TemporaryDirectory()
-        file = os.path.join(tdir.name, "model.npz")
-        m.save_weights(file)
+        npz_file = os.path.join(tdir.name, "model.npz")
+        m.save_weights(npz_file)
         m_load = make_model()
-        m_load.load_weights(file)
+        m_load.load_weights(npz_file)
+        tdir.cleanup()
+
+        eq_tree = tree_map(mx.array_equal, m.parameters(), m_load.parameters())
+        self.assertTrue(all(tree_flatten(eq_tree)))
+
+    def test_save_safetensors_weights(self):
+        def make_model():
+            return nn.Sequential(nn.Linear(2, 2), nn.ReLU(), nn.Linear(2, 2))
+
+        m = make_model()
+        tdir = tempfile.TemporaryDirectory()
+        safetensors_file = os.path.join(tdir.name, "model.safetensors")
+        m.save_weights(safetensors_file)
+        m_load = make_model()
+        m_load.load_weights(safetensors_file)
         tdir.cleanup()
 
         eq_tree = tree_map(mx.array_equal, m.parameters(), m_load.parameters())
@@ -121,20 +136,20 @@ class TestLayers(mlx_tests.MLXTestCase):
         inputs = mx.zeros((10, 4))
         layer = nn.Identity()
         outputs = layer(inputs)
-        self.assertEqual(tuple(inputs.shape), tuple(outputs.shape))
+        self.assertEqual(inputs.shape, outputs.shape)
 
     def test_linear(self):
         inputs = mx.zeros((10, 4))
         layer = nn.Linear(input_dims=4, output_dims=8)
         outputs = layer(inputs)
-        self.assertEqual(tuple(outputs.shape), (10, 8))
+        self.assertEqual(outputs.shape, (10, 8))
 
     def test_bilinear(self):
         inputs1 = mx.zeros((10, 2))
         inputs2 = mx.zeros((10, 4))
         layer = nn.Bilinear(input1_dims=2, input2_dims=4, output_dims=6)
         outputs = layer(inputs1, inputs2)
-        self.assertEqual(tuple(outputs.shape), (10, 6))
+        self.assertEqual(outputs.shape, (10, 6))
 
     def test_group_norm(self):
         x = mx.arange(100, dtype=mx.float32)
@@ -558,12 +573,12 @@ class TestLayers(mlx_tests.MLXTestCase):
         c = nn.Conv1d(in_channels=C_in, out_channels=C_out, kernel_size=ks)
         c.weight = mx.ones_like(c.weight)
         y = c(x)
-        self.assertEqual(y.shape, [N, L - ks + 1, C_out])
+        self.assertEqual(y.shape, (N, L - ks + 1, C_out))
         self.assertTrue(mx.allclose(y, mx.full(y.shape, ks * C_in, mx.float32)))
 
         c = nn.Conv1d(in_channels=C_in, out_channels=C_out, kernel_size=ks, stride=2)
         y = c(x)
-        self.assertEqual(y.shape, [N, (L - ks + 1) // 2, C_out])
+        self.assertEqual(y.shape, (N, (L - ks + 1) // 2, C_out))
         self.assertTrue("bias" in c.parameters())
 
         c = nn.Conv1d(in_channels=C_in, out_channels=C_out, kernel_size=ks, bias=False)
@@ -573,7 +588,7 @@ class TestLayers(mlx_tests.MLXTestCase):
         x = mx.ones((4, 8, 8, 3))
         c = nn.Conv2d(3, 1, 8)
         y = c(x)
-        self.assertEqual(y.shape, [4, 1, 1, 1])
+        self.assertEqual(y.shape, (4, 1, 1, 1))
         c.weight = mx.ones_like(c.weight) / 8 / 8 / 3
         y = c(x)
         self.assertTrue(np.allclose(y[:, 0, 0, 0], x.mean(axis=(1, 2, 3))))
@@ -581,13 +596,13 @@ class TestLayers(mlx_tests.MLXTestCase):
         # 3x3 conv no padding stride 1
         c = nn.Conv2d(3, 8, 3)
         y = c(x)
-        self.assertEqual(y.shape, [4, 6, 6, 8])
+        self.assertEqual(y.shape, (4, 6, 6, 8))
         self.assertLess(mx.abs(y - c.weight.sum((1, 2, 3))).max(), 1e-4)
 
         # 3x3 conv padding 1 stride 1
         c = nn.Conv2d(3, 8, 3, padding=1)
         y = c(x)
-        self.assertEqual(y.shape, [4, 8, 8, 8])
+        self.assertEqual(y.shape, (4, 8, 8, 8))
         self.assertLess(mx.abs(y[:, 1:7, 1:7] - c.weight.sum((1, 2, 3))).max(), 1e-4)
         self.assertLess(
             mx.abs(y[:, 0, 0] - c.weight[:, 1:, 1:].sum(axis=(1, 2, 3))).max(),
@@ -609,14 +624,14 @@ class TestLayers(mlx_tests.MLXTestCase):
         # 3x3 conv no padding stride 2
         c = nn.Conv2d(3, 8, 3, padding=0, stride=2)
         y = c(x)
-        self.assertEqual(y.shape, [4, 3, 3, 8])
+        self.assertEqual(y.shape, (4, 3, 3, 8))
         self.assertLess(mx.abs(y - c.weight.sum((1, 2, 3))).max(), 1e-4)
 
     def test_sequential(self):
         x = mx.ones((10, 2))
         m = nn.Sequential(nn.Linear(2, 10), nn.ReLU(), nn.Linear(10, 1))
         y = m(x)
-        self.assertEqual(y.shape, [10, 1])
+        self.assertEqual(y.shape, (10, 1))
         params = m.parameters()
         self.assertTrue("layers" in params)
         self.assertEqual(len(params["layers"]), 3)
@@ -652,7 +667,7 @@ class TestLayers(mlx_tests.MLXTestCase):
         x = mx.arange(10)
         y = m(x)
 
-        self.assertEqual(y.shape, [10, 16])
+        self.assertEqual(y.shape, (10, 16))
         similarities = y @ y.T
         self.assertLess(
             mx.abs(similarities[mx.arange(10), mx.arange(10)] - 1).max(), 1e-5
@@ -671,19 +686,19 @@ class TestLayers(mlx_tests.MLXTestCase):
         x = mx.array([1.0, -1.0, 0.0])
         y = nn.relu(x)
         self.assertTrue(mx.array_equal(y, mx.array([1.0, 0.0, 0.0])))
-        self.assertEqual(y.shape, [3])
+        self.assertEqual(y.shape, (3,))
         self.assertEqual(y.dtype, mx.float32)
 
     def test_leaky_relu(self):
         x = mx.array([1.0, -1.0, 0.0])
         y = nn.leaky_relu(x)
         self.assertTrue(mx.array_equal(y, mx.array([1.0, -0.01, 0.0])))
-        self.assertEqual(y.shape, [3])
+        self.assertEqual(y.shape, (3,))
         self.assertEqual(y.dtype, mx.float32)
 
         y = nn.LeakyReLU(negative_slope=0.1)(x)
         self.assertTrue(mx.array_equal(y, mx.array([1.0, -0.1, 0.0])))
-        self.assertEqual(y.shape, [3])
+        self.assertEqual(y.shape, (3,))
         self.assertEqual(y.dtype, mx.float32)
 
     def test_elu(self):
@@ -692,21 +707,21 @@ class TestLayers(mlx_tests.MLXTestCase):
         epsilon = 1e-4
         expected_y = mx.array([1.0, -0.6321, 0.0])
         self.assertTrue(mx.all(mx.abs(y - expected_y) < epsilon))
-        self.assertEqual(y.shape, [3])
+        self.assertEqual(y.shape, (3,))
         self.assertEqual(y.dtype, mx.float32)
 
         y = nn.ELU(alpha=1.1)(x)
         epsilon = 1e-4
         expected_y = mx.array([1.0, -0.6953, 0.0])
         self.assertTrue(mx.all(mx.abs(y - expected_y) < epsilon))
-        self.assertEqual(y.shape, [3])
+        self.assertEqual(y.shape, (3,))
         self.assertEqual(y.dtype, mx.float32)
 
     def test_relu6(self):
         x = mx.array([1.0, -1.0, 0.0, 7.0, -7.0])
         y = nn.relu6(x)
         self.assertTrue(mx.array_equal(y, mx.array([1.0, 0.0, 0.0, 6.0, 0.0])))
-        self.assertEqual(y.shape, [5])
+        self.assertEqual(y.shape, (5,))
         self.assertEqual(y.dtype, mx.float32)
 
     def test_softmax(self):
@@ -715,7 +730,7 @@ class TestLayers(mlx_tests.MLXTestCase):
         epsilon = 1e-4
         expected_y = mx.array([0.6652, 0.0900, 0.2447])
         self.assertTrue(mx.all(mx.abs(y - expected_y) < epsilon))
-        self.assertEqual(y.shape, [3])
+        self.assertEqual(y.shape, (3,))
         self.assertEqual(y.dtype, mx.float32)
 
     def test_softplus(self):
@@ -724,7 +739,7 @@ class TestLayers(mlx_tests.MLXTestCase):
         epsilon = 1e-4
         expected_y = mx.array([1.3133, 0.3133, 0.6931])
         self.assertTrue(mx.all(mx.abs(y - expected_y) < epsilon))
-        self.assertEqual(y.shape, [3])
+        self.assertEqual(y.shape, (3,))
         self.assertEqual(y.dtype, mx.float32)
 
     def test_softsign(self):
@@ -733,7 +748,22 @@ class TestLayers(mlx_tests.MLXTestCase):
         epsilon = 1e-4
         expected_y = mx.array([0.5, -0.5, 0.0])
         self.assertTrue(mx.all(mx.abs(y - expected_y) < epsilon))
-        self.assertEqual(y.shape, [3])
+        self.assertEqual(y.shape, (3,))
+        self.assertEqual(y.dtype, mx.float32)
+
+    def test_softshrink(self):
+        x = mx.array([1.0, -1.0, 0.0])
+        y = nn.softshrink(x)
+        epsilon = 1e-4
+        expected_y = mx.array([0.5, -0.5, 0.0])
+        self.assertTrue(mx.all(mx.abs(y - expected_y) < epsilon))
+        self.assertEqual(y.shape, (3,))
+        self.assertEqual(y.dtype, mx.float32)
+
+        y = nn.Softshrink(lambd=0.7)(x)
+        expected_y = mx.array([0.3, -0.3, 0.0])
+        self.assertTrue(mx.all(mx.abs(y - expected_y) < epsilon))
+        self.assertEqual(y.shape, (3,))
         self.assertEqual(y.dtype, mx.float32)
 
     def test_celu(self):
@@ -742,13 +772,13 @@ class TestLayers(mlx_tests.MLXTestCase):
         epsilon = 1e-4
         expected_y = mx.array([1.0, -0.6321, 0.0])
         self.assertTrue(mx.all(mx.abs(y - expected_y) < epsilon))
-        self.assertEqual(y.shape, [3])
+        self.assertEqual(y.shape, (3,))
         self.assertEqual(y.dtype, mx.float32)
 
         y = nn.CELU(alpha=1.1)(x)
         expected_y = mx.array([1.0, -0.6568, 0.0])
         self.assertTrue(mx.all(mx.abs(y - expected_y) < epsilon))
-        self.assertEqual(y.shape, [3])
+        self.assertEqual(y.shape, (3,))
         self.assertEqual(y.dtype, mx.float32)
 
     def test_log_softmax(self):
@@ -757,7 +787,7 @@ class TestLayers(mlx_tests.MLXTestCase):
         epsilon = 1e-4
         expected_y = mx.array([-2.4076, -1.4076, -0.4076])
         self.assertTrue(mx.all(mx.abs(y - expected_y) < epsilon))
-        self.assertEqual(y.shape, [3])
+        self.assertEqual(y.shape, (3,))
         self.assertEqual(y.dtype, mx.float32)
 
     def test_log_sigmoid(self):
@@ -766,7 +796,7 @@ class TestLayers(mlx_tests.MLXTestCase):
         epsilon = 1e-4
         expected_y = mx.array([-0.3133, -1.3133, -0.6931])
         self.assertTrue(mx.all(mx.abs(y - expected_y) < epsilon))
-        self.assertEqual(y.shape, [3])
+        self.assertEqual(y.shape, (3,))
         self.assertEqual(y.dtype, mx.float32)
 
     def test_prelu(self):
@@ -787,7 +817,7 @@ class TestLayers(mlx_tests.MLXTestCase):
         epsilon = 1e-4
         expected_y = mx.array([0.0, -0.375, 0.0, 1.125, 3.0])
         self.assertTrue(mx.all(mx.abs(y - expected_y) < epsilon))
-        self.assertEqual(y.shape, [5])
+        self.assertEqual(y.shape, (5,))
         self.assertEqual(y.dtype, mx.float32)
 
     def test_glu(self):

python/tests/test_ops.py

@@ -12,12 +12,12 @@ import numpy as np
 class TestOps(mlx_tests.MLXTestCase):
     def test_full_ones_zeros(self):
         x = mx.full(2, 3.0)
-        self.assertEqual(x.shape, [2])
+        self.assertEqual(x.shape, (2,))
         self.assertEqual(x.tolist(), [3.0, 3.0])
 
         x = mx.full((2, 3), 2.0)
         self.assertEqual(x.dtype, mx.float32)
-        self.assertEqual(x.shape, [2, 3])
+        self.assertEqual(x.shape, (2, 3))
         self.assertEqual(x.tolist(), [[2, 2, 2], [2, 2, 2]])
 
         x = mx.full([3, 2], mx.array([False, True]))
@@ -28,11 +28,11 @@ class TestOps(mlx_tests.MLXTestCase):
         self.assertEqual(x.tolist(), [[2, 3], [2, 3], [2, 3]])
 
         x = mx.zeros(2)
-        self.assertEqual(x.shape, [2])
+        self.assertEqual(x.shape, (2,))
         self.assertEqual(x.tolist(), [0.0, 0.0])
 
         x = mx.ones(2)
-        self.assertEqual(x.shape, [2])
+        self.assertEqual(x.shape, (2,))
         self.assertEqual(x.tolist(), [1.0, 1.0])
 
         for t in [mx.bool_, mx.int32, mx.float32]:
@@ -351,6 +351,14 @@ class TestOps(mlx_tests.MLXTestCase):
 
         self.assertEqual(mx.isinf(0 * mx.array(float("inf"))).tolist(), False)
 
+        x = mx.array([-2147483648, 0, 2147483647], dtype=mx.int32)
+        result = mx.isinf(x)
+        self.assertEqual(result.tolist(), [False, False, False])
+
+        x = mx.array([-32768, 0, 32767], dtype=mx.int16)
+        result = mx.isinf(x)
+        self.assertEqual(result.tolist(), [False, False, False])
+
     def test_tri(self):
         for shape in [[4], [4, 4], [2, 10]]:
             for diag in [-1, 0, 1, -2]:
@@ -378,6 +386,11 @@ class TestOps(mlx_tests.MLXTestCase):
         expected = [0, -7, 3]
         self.assertListEqual(mx.minimum(x, y).tolist(), expected)
 
+        a = mx.array([float("nan")])
+        b = mx.array([0.0])
+        self.assertTrue(math.isnan(mx.minimum(a, b).item()))
+        self.assertTrue(math.isnan(mx.minimum(b, a).item()))
+
     def test_maximum(self):
         x = mx.array([0.0, -5, 10.0])
         y = mx.array([1.0, -7.0, 3.0])
@@ -385,6 +398,11 @@ class TestOps(mlx_tests.MLXTestCase):
         expected = [1, -5, 10]
         self.assertListEqual(mx.maximum(x, y).tolist(), expected)
 
+        a = mx.array([float("nan")])
+        b = mx.array([0.0])
+        self.assertTrue(math.isnan(mx.maximum(a, b).item()))
+        self.assertTrue(math.isnan(mx.maximum(b, a).item()))
+
     def test_floor(self):
         x = mx.array([-22.03, 19.98, -27, 9, 0.0, -np.inf, np.inf])
         expected = [-23, 19, -27, 9, 0, -np.inf, np.inf]
@@ -416,6 +434,14 @@ class TestOps(mlx_tests.MLXTestCase):
 
         self.assertEqual(mx.isposinf(0 * mx.array(float("inf"))).tolist(), False)
 
+        x = mx.array([-2147483648, 0, 2147483647], dtype=mx.int32)
+        result = mx.isposinf(x)
+        self.assertEqual(result.tolist(), [False, False, False])
+
+        x = mx.array([-32768, 0, 32767], dtype=mx.int16)
+        result = mx.isposinf(x)
+        self.assertEqual(result.tolist(), [False, False, False])
+
     def test_isneginf(self):
         x = mx.array([0.0, float("-inf")])
         self.assertEqual(mx.isneginf(x).tolist(), [False, True])
@@ -431,6 +457,14 @@ class TestOps(mlx_tests.MLXTestCase):
 
         self.assertEqual(mx.isneginf(0 * mx.array(float("inf"))).tolist(), False)
 
+        x = mx.array([-2147483648, 0, 2147483647], dtype=mx.int32)
+        result = mx.isneginf(x)
+        self.assertEqual(result.tolist(), [False, False, False])
+
+        x = mx.array([-32768, 0, 32767], dtype=mx.int16)
+        result = mx.isneginf(x)
+        self.assertEqual(result.tolist(), [False, False, False])
+
     def test_round(self):
         # float
         x = mx.array(
@@ -496,10 +530,10 @@ class TestOps(mlx_tests.MLXTestCase):
 
     def test_move_swap_axes(self):
         x = mx.zeros((2, 3, 4))
-        self.assertEqual(mx.moveaxis(x, 0, 2).shape, [3, 4, 2])
-        self.assertEqual(x.moveaxis(0, 2).shape, [3, 4, 2])
-        self.assertEqual(mx.swapaxes(x, 0, 2).shape, [4, 3, 2])
-        self.assertEqual(x.swapaxes(0, 2).shape, [4, 3, 2])
+        self.assertEqual(mx.moveaxis(x, 0, 2).shape, (3, 4, 2))
+        self.assertEqual(x.moveaxis(0, 2).shape, (3, 4, 2))
+        self.assertEqual(mx.swapaxes(x, 0, 2).shape, (4, 3, 2))
+        self.assertEqual(x.swapaxes(0, 2).shape, (4, 3, 2))
 
     def test_sum(self):
         x = mx.array(
@@ -511,7 +545,7 @@ class TestOps(mlx_tests.MLXTestCase):
         self.assertEqual(mx.sum(x).item(), 9)
         y = mx.sum(x, keepdims=True)
         self.assertEqual(y, mx.array(9))
-        self.assertEqual(y.shape, [1, 1])
+        self.assertEqual(y.shape, (1, 1))
 
         self.assertEqual(mx.sum(x, axis=0).tolist(), [4, 5])
         self.assertEqual(mx.sum(x, axis=1).tolist(), [3, 6])
@@ -551,7 +585,7 @@ class TestOps(mlx_tests.MLXTestCase):
         self.assertEqual(mx.prod(x).item(), 18)
         y = mx.prod(x, keepdims=True)
         self.assertEqual(y, mx.array(18))
-        self.assertEqual(y.shape, [1, 1])
+        self.assertEqual(y.shape, (1, 1))
 
         self.assertEqual(mx.prod(x, axis=0).tolist(), [3, 6])
         self.assertEqual(mx.prod(x, axis=1).tolist(), [2, 9])
@@ -566,11 +600,11 @@ class TestOps(mlx_tests.MLXTestCase):
         self.assertEqual(mx.min(x).item(), 1)
         self.assertEqual(mx.max(x).item(), 4)
         y = mx.min(x, keepdims=True)
-        self.assertEqual(y.shape, [1, 1])
+        self.assertEqual(y.shape, (1, 1))
         self.assertEqual(y, mx.array(1))
 
         y = mx.max(x, keepdims=True)
-        self.assertEqual(y.shape, [1, 1])
+        self.assertEqual(y.shape, (1, 1))
         self.assertEqual(y, mx.array(4))
 
         self.assertEqual(mx.min(x, axis=0).tolist(), [1, 2])
@@ -636,7 +670,7 @@ class TestOps(mlx_tests.MLXTestCase):
         self.assertEqual(mx.mean(x).item(), 2.5)
         y = mx.mean(x, keepdims=True)
         self.assertEqual(y, mx.array(2.5))
-        self.assertEqual(y.shape, [1, 1])
+        self.assertEqual(y.shape, (1, 1))
 
         self.assertEqual(mx.mean(x, axis=0).tolist(), [2, 3])
         self.assertEqual(mx.mean(x, axis=1).tolist(), [1.5, 3.5])
@@ -651,7 +685,7 @@ class TestOps(mlx_tests.MLXTestCase):
         self.assertEqual(mx.var(x).item(), 1.25)
         y = mx.var(x, keepdims=True)
         self.assertEqual(y, mx.array(1.25))
-        self.assertEqual(y.shape, [1, 1])
+        self.assertEqual(y.shape, (1, 1))
 
         self.assertEqual(mx.var(x, axis=0).tolist(), [1.0, 1.0])
         self.assertEqual(mx.var(x, axis=1).tolist(), [0.25, 0.25])
@@ -736,6 +770,10 @@ class TestOps(mlx_tests.MLXTestCase):
 
         self.assertTrue(np.allclose(result, expected))
 
+        a = mx.array([float("nan")])
+        b = mx.array([0.0])
+        self.assertTrue(math.isnan(mx.logaddexp(a, b).item()))
+
     def test_log(self):
         a = mx.array([1, 0.5, 10, 100])
         result = mx.log(a)
@@ -831,11 +869,26 @@ class TestOps(mlx_tests.MLXTestCase):
         self.assertFalse(mx.allclose(a, b, 0.01).item())
         self.assertTrue(mx.allclose(a, b, 0.01, 0.1).item())
 
+        c = mx.array(float("inf"))
+        self.assertTrue(mx.allclose(c, c).item())
+
+    def test_isclose(self):
+        a = mx.array([float("inf"), float("inf"), float("-inf")])
+        b = mx.array([float("inf"), float("-inf"), float("-inf")])
+
+        self.assertListEqual(mx.isclose(a, b).tolist(), [True, False, True])
+
+        a = mx.array([np.nan])
+        self.assertListEqual(mx.isclose(a, a).tolist(), [False])
+
+        a = mx.array([np.nan])
+        self.assertListEqual(mx.isclose(a, a, equal_nan=True).tolist(), [True])
+
     def test_all(self):
         a = mx.array([[True, False], [True, True]])
 
         self.assertFalse(mx.all(a).item())
-        self.assertEqual(mx.all(a, keepdims=True).shape, [1, 1])
+        self.assertEqual(mx.all(a, keepdims=True).shape, (1, 1))
         self.assertFalse(mx.all(a, axis=[0, 1]).item())
         self.assertEqual(mx.all(a, axis=[0]).tolist(), [True, False])
         self.assertEqual(mx.all(a, axis=[1]).tolist(), [False, True])
@@ -846,7 +899,7 @@ class TestOps(mlx_tests.MLXTestCase):
         a = mx.array([[True, False], [False, False]])
 
         self.assertTrue(mx.any(a).item())
-        self.assertEqual(mx.any(a, keepdims=True).shape, [1, 1])
+        self.assertEqual(mx.any(a, keepdims=True).shape, (1, 1))
         self.assertTrue(mx.any(a, axis=[0, 1]).item())
         self.assertEqual(mx.any(a, axis=[0]).tolist(), [True, False])
         self.assertEqual(mx.any(a, axis=[1]).tolist(), [True, False])
@@ -903,22 +956,22 @@ class TestOps(mlx_tests.MLXTestCase):
 
         a_npy_taken = np.take(a_npy, idx_npy)
         a_mlx_taken = mx.take(a_mlx, idx_mlx)
-        self.assertListEqual(list(a_npy_taken.shape), a_mlx_taken.shape)
+        self.assertEqual(a_npy_taken.shape, a_mlx_taken.shape)
         self.assertListEqual(a_npy_taken.tolist(), a_mlx_taken.tolist())
 
         a_npy_taken = np.take(a_npy, idx_npy, axis=0)
         a_mlx_taken = mx.take(a_mlx, idx_mlx, axis=0)
-        self.assertListEqual(list(a_npy_taken.shape), a_mlx_taken.shape)
+        self.assertEqual(a_npy_taken.shape, a_mlx_taken.shape)
         self.assertListEqual(a_npy_taken.tolist(), a_mlx_taken.tolist())
 
         a_npy_taken = np.take(a_npy, idx_npy, axis=1)
         a_mlx_taken = mx.take(a_mlx, idx_mlx, axis=1)
-        self.assertListEqual(list(a_npy_taken.shape), a_mlx_taken.shape)
+        self.assertEqual(a_npy_taken.shape, a_mlx_taken.shape)
         self.assertListEqual(a_npy_taken.tolist(), a_mlx_taken.tolist())
 
         a_npy_taken = np.take(a_npy, idx_npy, axis=2)
         a_mlx_taken = mx.take(a_mlx, idx_mlx, axis=2)
-        self.assertListEqual(list(a_npy_taken.shape), a_mlx_taken.shape)
+        self.assertEqual(a_npy_taken.shape, a_mlx_taken.shape)
         self.assertListEqual(a_npy_taken.tolist(), a_mlx_taken.tolist())
 
     def test_take_along_axis(self):
@@ -956,6 +1009,17 @@ class TestOps(mlx_tests.MLXTestCase):
         self.assertEqual(z.tolist(), [5, 6, 7])
 
     def test_arange_overload_dispatch(self):
+        with self.assertRaises(ValueError):
+            a = mx.arange(float("nan"), 1, 5)
+        with self.assertRaises(ValueError):
+            a = mx.arange(0, float("nan"), 5)
+        with self.assertRaises(ValueError):
+            a = mx.arange(0, 2, float("nan"))
+        with self.assertRaises(ValueError):
+            a = mx.arange(0, float("inf"), float("inf"))
+        with self.assertRaises(ValueError):
+            a = mx.arange(float("inf"), 1, float("inf"))
+
         a = mx.arange(5)
         expected = [0, 1, 2, 3, 4]
         self.assertListEqual(a.tolist(), expected)
@@ -1310,6 +1374,11 @@ class TestOps(mlx_tests.MLXTestCase):
                 self.assertEqual(list(c_npy.shape), list(c_mlx.shape))
                 self.assertTrue(np.allclose(c_npy, c_mlx, atol=1e-6))
 
+        with self.assertRaises(ValueError):
+            a = mx.array([[1, 2], [1, 2], [1, 2]])
+            b = mx.array([1, 2])
+            mx.concatenate([a, b], axis=0)
+
     def test_pad(self):
         pad_width_and_values = [
             ([(1, 1), (1, 1), (1, 1)], 0),
@@ -1331,13 +1400,13 @@ class TestOps(mlx_tests.MLXTestCase):
                 self.assertTrue(np.allclose(b_npy, b_mlx, atol=1e-6))
 
         a = mx.zeros((1, 1, 1))
-        self.assertEqual(mx.pad(a, 1).shape, [3, 3, 3])
-        self.assertEqual(mx.pad(a, (1,)).shape, [3, 3, 3])
-        self.assertEqual(mx.pad(a, [1]).shape, [3, 3, 3])
-        self.assertEqual(mx.pad(a, (1, 2)).shape, [4, 4, 4])
-        self.assertEqual(mx.pad(a, [(1, 2)]).shape, [4, 4, 4])
-        self.assertEqual(mx.pad(a, ((1, 2),)).shape, [4, 4, 4])
-        self.assertEqual(mx.pad(a, ((1, 2), (2, 1), (2, 2))).shape, [4, 4, 5])
+        self.assertEqual(mx.pad(a, 1).shape, (3, 3, 3))
+        self.assertEqual(mx.pad(a, (1,)).shape, (3, 3, 3))
+        self.assertEqual(mx.pad(a, [1]).shape, (3, 3, 3))
+        self.assertEqual(mx.pad(a, (1, 2)).shape, (4, 4, 4))
+        self.assertEqual(mx.pad(a, [(1, 2)]).shape, (4, 4, 4))
+        self.assertEqual(mx.pad(a, ((1, 2),)).shape, (4, 4, 4))
+        self.assertEqual(mx.pad(a, ((1, 2), (2, 1), (2, 2))).shape, (4, 4, 5))
 
         # Test grads
         a_fwd = mx.array(np.random.rand(16, 16).astype(np.float32))
@@ -1421,19 +1490,19 @@ class TestOps(mlx_tests.MLXTestCase):
 
     def test_squeeze_expand(self):
         a = mx.zeros((2, 1, 2, 1))
-        self.assertEqual(mx.squeeze(a).shape, [2, 2])
-        self.assertEqual(mx.squeeze(a, 1).shape, [2, 2, 1])
-        self.assertEqual(mx.squeeze(a, [1, 3]).shape, [2, 2])
-        self.assertEqual(a.squeeze().shape, [2, 2])
-        self.assertEqual(a.squeeze(1).shape, [2, 2, 1])
-        self.assertEqual(a.squeeze([1, 3]).shape, [2, 2])
+        self.assertEqual(mx.squeeze(a).shape, (2, 2))
+        self.assertEqual(mx.squeeze(a, 1).shape, (2, 2, 1))
+        self.assertEqual(mx.squeeze(a, [1, 3]).shape, (2, 2))
+        self.assertEqual(a.squeeze().shape, (2, 2))
+        self.assertEqual(a.squeeze(1).shape, (2, 2, 1))
+        self.assertEqual(a.squeeze([1, 3]).shape, (2, 2))
 
         a = mx.zeros((2, 2))
-        self.assertEqual(mx.squeeze(a).shape, [2, 2])
+        self.assertEqual(mx.squeeze(a).shape, (2, 2))
 
-        self.assertEqual(mx.expand_dims(a, 0).shape, [1, 2, 2])
-        self.assertEqual(mx.expand_dims(a, (0, 1)).shape, [1, 1, 2, 2])
-        self.assertEqual(mx.expand_dims(a, [0, -1]).shape, [1, 2, 2, 1])
+        self.assertEqual(mx.expand_dims(a, 0).shape, (1, 2, 2))
+        self.assertEqual(mx.expand_dims(a, (0, 1)).shape, (1, 1, 2, 2))
+        self.assertEqual(mx.expand_dims(a, [0, -1]).shape, (1, 2, 2, 1))
 
     def test_sort(self):
         shape = (3, 4, 5)
@@ -1534,12 +1603,12 @@ class TestOps(mlx_tests.MLXTestCase):
 
     def test_flatten(self):
         x = mx.zeros([2, 3, 4])
-        self.assertEqual(mx.flatten(x).shape, [2 * 3 * 4])
-        self.assertEqual(mx.flatten(x, start_axis=1).shape, [2, 3 * 4])
-        self.assertEqual(mx.flatten(x, end_axis=1).shape, [2 * 3, 4])
-        self.assertEqual(x.flatten().shape, [2 * 3 * 4])
-        self.assertEqual(x.flatten(start_axis=1).shape, [2, 3 * 4])
-        self.assertEqual(x.flatten(end_axis=1).shape, [2 * 3, 4])
+        self.assertEqual(mx.flatten(x).shape, (2 * 3 * 4,))
+        self.assertEqual(mx.flatten(x, start_axis=1).shape, (2, 3 * 4))
+        self.assertEqual(mx.flatten(x, end_axis=1).shape, (2 * 3, 4))
+        self.assertEqual(x.flatten().shape, (2 * 3 * 4,))
+        self.assertEqual(x.flatten(start_axis=1).shape, (2, 3 * 4))
+        self.assertEqual(x.flatten(end_axis=1).shape, (2 * 3, 4))
 
     def test_clip(self):
         a = np.array([1, 4, 3, 8, 5], np.int32)
@@ -1620,23 +1689,23 @@ class TestOps(mlx_tests.MLXTestCase):
                 self.assertCmpNumpy(
                     [(3, 4, 5), (4, 3, 2)],
                     mx.tensordot,
-                    lambda x, y, dims: np.tensordot(x, y, axes=dims),
+                    np.tensordot,
                     dtype=dtype,
-                    dims=([1, 0], [0, 1]),
+                    axes=([1, 0], [0, 1]),
                 )
                 self.assertCmpNumpy(
                     [(3, 4, 5), (4, 5, 6)],
                     mx.tensordot,
-                    lambda x, y, dims: np.tensordot(x, y, axes=dims),
+                    np.tensordot,
                     dtype=dtype,
-                    dims=2,
+                    axes=2,
                 )
                 self.assertCmpNumpy(
                     [(3, 5, 4, 6), (6, 4, 5, 3)],
                     mx.tensordot,
-                    lambda x, y, dims: np.tensordot(x, y, axes=dims),
+                    np.tensordot,
                     dtype=dtype,
-                    dims=([2, 1, 3], [1, 2, 0]),
+                    axes=([2, 1, 3], [1, 2, 0]),
                 )
 
     def test_inner(self):
@@ -1706,6 +1775,72 @@ class TestOps(mlx_tests.MLXTestCase):
         )
         self.assertCmpNumpy([(3,), [2, 2, 2]], mx.tile, np.tile)
 
+    def test_empty_matmuls(self):
+        a = mx.array([])
+        b = mx.array([])
+        self.assertEqual(mx.inner(a, b).item(), 0.0)
+
+        a = mx.zeros((10, 0))
+        b = mx.zeros((0, 10))
+        out = a @ b
+        self.assertTrue(mx.array_equal(out, mx.zeros((10, 10))))
+
+    def test_diagonal(self):
+        x = mx.array(
+            [
+                [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]],
+                [[12, 13, 14, 15], [16, 17, 18, 19], [20, 21, 22, 23]],
+            ]
+        )
+        expected = [[0, 13], [4, 17], [8, 21]]
+
+        self.assertListEqual(mx.diagonal(x, 0, -1, 0).tolist(), expected)
+
+        expected = [[1, 14], [5, 18], [9, 22]]
+        self.assertListEqual(mx.diagonal(x, -1, 2, 0).tolist(), expected)
+
+    def test_diag(self):
+        # Test 1D input
+        x = mx.array([1, 2, 3, 4])
+        expected = mx.array([[1, 0, 0, 0], [0, 2, 0, 0], [0, 0, 3, 0], [0, 0, 0, 4]])
+        result = mx.diag(x)
+        self.assertTrue(mx.array_equal(result, expected))
+
+        # Test 1D with offset
+        x = mx.array([2, 6])
+        result = mx.diag(x, k=5)
+        expected = mx.array(np.diag(x, k=5))
+        self.assertTrue(mx.array_equal(result, expected))
+
+        # Test 2D input
+        x = mx.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
+        expected = mx.array([1, 5, 9])
+        result = mx.diag(x)
+        self.assertTrue(mx.array_equal(result, expected))
+
+        # Test with offset
+        expected = mx.array([2, 6])
+        result = mx.diag(x, 1)
+        self.assertTrue(mx.array_equal(result, expected))
+
+        # Test non-square
+        x = mx.array([[1, 2, 3], [4, 5, 6]])
+        result = mx.diag(x)
+        expected = mx.array(np.diag(x))
+        self.assertTrue(mx.array_equal(result, expected))
+
+        result = mx.diag(x, k=10)
+        expected = mx.array(np.diag(x, k=10))
+        self.assertTrue(mx.array_equal(result, expected))
+
+        result = mx.diag(x, k=-10)
+        expected = mx.array(np.diag(x, k=-10))
+        self.assertTrue(mx.array_equal(result, expected))
+
+        result = mx.diag(x, k=-1)
+        expected = mx.array(np.diag(x, k=-1))
+        self.assertTrue(mx.array_equal(result, expected))
+
 
 if __name__ == "__main__":
     unittest.main()

python/tests/test_optimizers.py

@@ -39,6 +39,24 @@ class TestOptimizers(mlx_tests.MLXTestCase):
             all_equal = all(v for _, v in mlx.utils.tree_flatten(equal_shape))
             self.assertTrue(all_equal)
 
+    def test_adafactor(self):
+        x = mx.zeros((5, 5))
+        grad = mx.ones_like(x)
+        optimizer = opt.Adafactor()
+        for _ in range(2):
+            xp = optimizer.apply_single(grad, x, optimizer.state)
+            self.assertEqual(xp.dtype, x.dtype)
+            self.assertEqual(xp.shape, x.shape)
+
+        x = mx.zeros((5, 5), mx.float16)
+        grad = mx.ones_like(x)
+        optimizer = opt.Adafactor()
+        for _ in range(2):
+            xp = optimizer.apply_single(grad, x, optimizer.state)
+            self.assertEqual(xp.dtype, x.dtype)
+            self.assertEqual(xp.shape, x.shape)
+        self.assertEqual(optimizer.state["step"], 2)
+
 
 if __name__ == "__main__":
     unittest.main()