Bumping the version (#256)

* Change the default quantization group_size to 64
* Rename groups to group_size and width to bits

.circleci/config.yml

@@ -203,7 +203,7 @@ workflows:
               ignore: /.*/
           matrix:
             parameters:
-              python_version: ["3.8", "3.9", "3.10", "3.11"]
+              python_version: ["3.8", "3.9", "3.10", "3.11", "3.12"]
               macos_version: ["13", "14"]
   nightly_build:
     when: << pipeline.parameters.nightly_build >>
@@ -211,7 +211,7 @@ workflows:
       - build_package:
           matrix:
             parameters:
-              python_version: ["3.8", "3.9", "3.10", "3.11"]
+              python_version: ["3.8", "3.9", "3.10", "3.11", "3.12"]
               macos_version: ["13", "14"]
   weekly_build:
     when: << pipeline.parameters.weekly_build >>
@@ -219,5 +219,5 @@ workflows:
       - build_dev_release:
           matrix:
             parameters:
-              python_version: ["3.8", "3.9", "3.10", "3.11"]
+              python_version: ["3.8", "3.9", "3.10", "3.11", "3.12"]
               macos_version: ["13", "14"]

.github/pull_request_template.md

@@ -0,0 +1,12 @@
+## Proposed changes
+
+Please include a description of the problem or feature this PR is addressing. If there is a corresponding issue, include the issue #.
+
+## Checklist
+
+Put an `x` in the boxes that apply.
+
+- [ ] I have read the [CONTRIBUTING](https://github.com/ml-explore/mlx/blob/main/CONTRIBUTING.md) document
+- [ ] I have run `pre-commit run --all-files` to format my code / installed pre-commit prior to committing changes
+- [ ] I have added tests that prove my fix is effective or that my feature works
+- [ ] I have updated the necessary documentation (if needed)

.github/workflows/pull_request.yml

@@ -0,0 +1,20 @@
+on:
+  pull_request:
+    branches:
+      - main
+
+jobs:
+  check_lint:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v4
+        with:
+          python-version: 3.8
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install pre-commit black isort clang-format
+      - name: Run lint
+        run: |
+          pre-commit run --all-files

.gitignore

@@ -8,8 +8,10 @@ __pycache__/
 
 # Metal libraries
 *.metallib
+venv/
 
 # Distribution / packaging
+python/mlx/core
 python/mlx/share
 python/mlx/include
 .Python
@@ -73,3 +75,6 @@ build/
 # VSCode 
 .vscode/
 .DS_Store
+
+# Jetbrains
+.cache

.pre-commit-config.yaml

@@ -1,9 +1,16 @@
 repos:
 -   repo: https://github.com/pre-commit/mirrors-clang-format
-    rev: v14.0.6
+    rev: v17.0.6
     hooks:
     -   id: clang-format
--   repo: https://github.com/psf/black
+# Using this mirror lets us use mypyc-compiled black, which is about 2x faster
+-   repo: https://github.com/psf/black-pre-commit-mirror
     rev: 22.10.0
     hooks:
     -   id: black
+-   repo: https://github.com/pycqa/isort
+    rev: 5.12.0
+    hooks:
+    -   id: isort
+        args:
+            - --profile=black

ACKNOWLEDGMENTS.md

@@ -1,3 +1,17 @@
+# Individual Contributors
+
+If you wish to be acknowledged for your contributions, please list your name
+with a short description of your contribution(s) below. For example:
+
+- Jane Smith: Added the `foo` and `bar` ops.
+
+MLX was developed with contributions from the following individuals:
+  
+- Juarez Bochi: Fixed bug in cross attention.
+- Justin Deschenaux: Sine, Cosine, arange, randint, truncated normal, bernoulli, lion optimizer, linear and logistic regression python example.
+
+# Third-Party Software
+
 MLX leverages several third-party software, listed here together with
 their license copied verbatim.
 
@@ -231,4 +245,4 @@ Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
-limitations under the License.
+limitations under the License.

CMakeLists.txt

@@ -18,7 +18,28 @@ 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.1)
+  set(MLX_VERSION 0.0.6)
+endif()
+
+# --------------------- Processor tests -------------------------
+
+message(STATUS "Building MLX for ${CMAKE_HOST_SYSTEM_PROCESSOR} processor on ${CMAKE_SYSTEM_NAME}")
+
+set(MLX_BUILD_ARM OFF)
+
+if (${CMAKE_SYSTEM_NAME} MATCHES "Darwin")
+
+  if (${CMAKE_HOST_SYSTEM_PROCESSOR} MATCHES "x86_64")
+    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")
+    set(MLX_BUILD_ARM ON)
+  endif()
+
+else()
+  message(WARNING "MLX is prioritised for Apple silicon systems using macOS.")
 endif()
 
 # ----------------------------- Lib -----------------------------
@@ -37,20 +58,26 @@ endif()
 
 if (MLX_BUILD_METAL AND NOT METAL_LIB)
   message(STATUS "Metal not found. Unable to build GPU")
+  set(MLX_BUILD_METAL OFF)
 elseif (MLX_BUILD_METAL)
   message(STATUS "Building METAL sources")
   add_compile_definitions(_METAL_)
 
-  execute_process(COMMAND zsh "-c" "/usr/bin/sw_vers | cut -f2- -d: | sed -n 2p | grep -Eo '[0-9]+.[0-9]+'"
-                  OUTPUT_VARIABLE MACOS_VERSION)
+  # Throw an error if xcrun not found
+  execute_process(COMMAND zsh "-c" "/usr/bin/xcrun -sdk macosx --show-sdk-version"
+                  OUTPUT_VARIABLE MACOS_VERSION
+                  COMMAND_ERROR_IS_FATAL ANY)
 
-  message(STATUS "Detected macOS version ${MACOS_VERSION}")
-  if (${MACOS_VERSION} GREATER_EQUAL 14.0)
+  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)
     set(METAL_CPP_URL https://developer.apple.com/metal/cpp/files/metal-cpp_macOS14_iOS17-beta.zip)
   elseif (${MACOS_VERSION} GREATER_EQUAL 13.3)
     set(METAL_CPP_URL https://developer.apple.com/metal/cpp/files/metal-cpp_macOS13.3_iOS16.4.zip)
   else()
-    set(METAL_CPP_URL https://developer.apple.com/metal/cpp/files/metal-cpp_macOS13_iOS16.zip)
+    message(FATAL_ERROR "MLX requires macOS >= 13.4 to be built with MLX_BUILD_METAL=ON" )
   endif()
 
   FetchContent_Declare(
@@ -72,13 +99,13 @@ elseif (MLX_BUILD_METAL)
 endif()
 
 find_library(ACCELERATE_LIBRARY Accelerate)
-if (ACCELERATE_LIBRARY)
+if (MLX_BUILD_ARM AND ACCELERATE_LIBRARY)
   message(STATUS "Accelerate found ${ACCELERATE_LIBRARY}")
   set(MLX_BUILD_ACCELERATE ON)
   target_link_libraries(mlx ${ACCELERATE_LIBRARY})
   add_compile_definitions(ACCELERATE_NEW_LAPACK)
 else()
-  message(STATUS "Accelerate not found, using default backend.")
+  message(STATUS "Accelerate or arm neon not found, using default backend.")
   set(MLX_BUILD_ACCELERATE OFF)
   #set(BLA_VENDOR Generic)
   find_package(BLAS REQUIRED)
@@ -194,4 +221,4 @@ install(
 install(
   DIRECTORY ${CMAKE_MODULE_PATH}/
   DESTINATION ${MLX_CMAKE_INSTALL_MODULE_DIR}
-)
+)

README.md

@@ -2,38 +2,41 @@
 
 [**Quickstart**](#quickstart) | [**Installation**](#installation) |
 [**Documentation**](https://ml-explore.github.io/mlx/build/html/index.html) |
-[**Examples**](#examples)
+[**Examples**](#examples) 
 
-MLX is an array framework for machine learning on Apple silicon.
+[![CircleCI](https://circleci.com/gh/ml-explore/mlx.svg?style=svg)](https://circleci.com/gh/ml-explore/mlx)
+
+MLX is an array framework for machine learning on Apple silicon, brought to you
+by Apple machine learning research.
 
 Some key features of MLX include:
 
- - **Familiar APIs**: MLX has a Python API which closely follows NumPy.
-   MLX also has a fully featured C++ API which closely mirrors the Python API. 
-   MLX has higher level packages like `mlx.nn` and `mlx.optimizers` with APIs
+ - **Familiar APIs**: MLX has a Python API that closely follows NumPy.
+   MLX also has a fully featured C++ API, which closely mirrors the Python API. 
+   MLX has higher-level packages like `mlx.nn` and `mlx.optimizers` with APIs
    that closely follow PyTorch to simplify building more complex models.
 
- - **Composable function transformations**: MLX has composable function
+ - **Composable function transformations**: MLX supports composable function
    transformations for automatic differentiation, automatic vectorization,
    and computation graph optimization.
 
  - **Lazy computation**: Computations in MLX are lazy. Arrays are only
    materialized when needed.
 
- - **Dynamic graph construction**: Computation graphs in MLX are built
+ - **Dynamic graph construction**: Computation graphs in MLX are constructed
    dynamically. Changing the shapes of function arguments does not trigger
    slow compilations, and debugging is simple and intuitive.
 
  - **Multi-device**: Operations can run on any of the supported devices
-   (currently the CPU and GPU).
+   (currently the CPU and the GPU).
 
- - **Unified memory**: A noteable difference from MLX and other frameworks
+ - **Unified memory**: A notable difference from MLX and other frameworks
    is the *unified memory model*. Arrays in MLX live in shared memory.
    Operations on MLX arrays can be performed on any of the supported
-   device types without moving data.
+   device types without transferring data.
 
 MLX is designed by machine learning researchers for machine learning
-researchers. The framework is intended to be user friendly, but still efficient
+researchers. The framework is intended to be user-friendly, but still efficient
 to train and deploy models. The design of the framework itself is also
 conceptually simple. We intend to make it easy for researchers to extend and
 improve MLX with the goal of quickly exploring new ideas. 
@@ -46,10 +49,10 @@ The design of MLX is inspired by frameworks like
 ## Examples
 
 The [MLX examples repo](https://github.com/ml-explore/mlx-examples) has a
-variety of examples including:
+variety of examples, including:
 
 - [Transformer language model](https://github.com/ml-explore/mlx-examples/tree/main/transformer_lm) training.
-- Large scale text generation with
+- Large-scale text generation with
   [LLaMA](https://github.com/ml-explore/mlx-examples/tree/main/llama) and
   finetuning with [LoRA](https://github.com/ml-explore/mlx-examples/tree/main/lora).
 - Generating images with [Stable Diffusion](https://github.com/ml-explore/mlx-examples/tree/main/stable_diffusion).
@@ -63,7 +66,7 @@ in the documentation.
 
 ## Installation
 
-MLX is available on [PyPi](https://pypi.org/project/mlx/). To install the Python API run:
+MLX is available on [PyPI](https://pypi.org/project/mlx/). To install the Python API, run:
 
 ```
 pip install mlx
@@ -76,4 +79,28 @@ for more information on building the C++ and Python APIs from source.
 ## Contributing 
 
 Check out the [contribution guidelines](CONTRIBUTING.md) for more information
-on contributing to MLX.
+on contributing to MLX. See the
+[docs](https://ml-explore.github.io/mlx/build/html/install.html) for more
+information on building from source, and running tests.
+
+We are grateful for all of [our
+contributors](ACKNOWLEDGMENTS.md#Individual-Contributors). If you contribute
+to MLX and wish to be acknowledged, please add your name to to the list in your
+pull request.
+
+## Citing MLX
+
+The MLX software suite was initially developed with equal contribution by Awni
+Hannun, Jagrit Digani, Angelos Katharopoulos, and Ronan Collobert. If you find
+MLX useful in your research and wish to cite it, please use the following
+BibTex entry:
+
+```
+@software{mlx2023,
+  author = {Awni Hannun and Jagrit Digani and Angelos Katharopoulos and Ronan Collobert},
+  title = {{MLX}: Efficient and flexible machine learning on Apple silicon},
+  url = {https://github.com/ml-explore},
+  version = {0.0},
+  year = {2023},
+}
+```

benchmarks/numpy/single_ops.py

@@ -1,7 +1,6 @@
 # Copyright © 2023 Apple Inc.
 
 import numpy as np
-
 from time_utils import time_fn
 
 

benchmarks/python/batch_matmul_bench.py

@@ -1,8 +1,8 @@
 # Copyright © 2023 Apple Inc.
 
 import argparse
-import mlx.core as mx
 
+import mlx.core as mx
 from time_utils import time_fn
 
 B = 8
@@ -30,7 +30,7 @@ def time_batch_matmul():
     time_fn(batch_vjp_second)
 
 
-def time_unbatch_matmul(key):
+def time_unbatch_matmul():
     mx.random.seed(3)
     a = mx.random.uniform(shape=(B * T, D))
     b = mx.random.uniform(shape=(D, D))

benchmarks/python/blas/bench_gemm.py

@@ -1,13 +1,14 @@
 # Copyright © 2023 Apple Inc.
 
-import numpy as np
 import argparse
-import mlx.core as mx
-import time
-import torch
-import os
 import math
+import os
 import subprocess
+import time
+
+import mlx.core as mx
+import numpy as np
+import torch
 
 device_name = subprocess.check_output(["sysctl", "-n", "machdep.cpu.brand_string"])
 device_name = device_name.decode("utf-8").strip("\n")

benchmarks/python/blas/bench_gemv.py

@@ -1,14 +1,14 @@
 # Copyright © 2023 Apple Inc.
 
-import matplotlib.pyplot as plt
-import numpy as np
 import argparse
-import mlx.core as mx
-import time
-import torch
 import os
 import subprocess
+import time
 
+import matplotlib.pyplot as plt
+import mlx.core as mx
+import numpy as np
+import torch
 
 results_dir = "./results"
 

benchmarks/python/comparative/bench_mlx.py

@@ -6,6 +6,7 @@ import os
 import time
 
 import mlx.core as mx
+import mlx.nn as nn
 
 
 def int_or_list(x):
@@ -22,6 +23,16 @@ def none_or_list(x):
         return [int(xi) for xi in x.split(",")]
 
 
+def dtype_from_str(x):
+    if x == "":
+        return mx.float32
+    else:
+        dt = getattr(mx, x)
+        if not isinstance(dt, mx.Dtype):
+            raise ValueError(f"{x} is not an mlx dtype")
+        return dt
+
+
 def bench(f, *args):
     for i in range(10):
         f(*args)
@@ -48,6 +59,15 @@ def matmul(x, y):
     mx.eval(ys)
 
 
+def quant_matmul(x, w, s, b):
+    groups = x.shape[-1] // s.shape[-1]
+    width = 32 // (x.shape[-1] // w.shape[0])
+    ys = []
+    for i in range(10):
+        ys.append(mx.quantized_matmul(x, w, s, b, groups=groups, width=width))
+    mx.eval(ys)
+
+
 def conv1d(x, y):
     ys = []
     for i in range(10):
@@ -95,7 +115,77 @@ def softmax_fused(axis, x):
 def relu(x):
     y = x
     for i in range(100):
-        y = mx.maximum(y, 0)
+        y = nn.relu(y)
+    mx.eval(y)
+
+
+def leaky_relu(x: mx.array):
+    y = x
+    for i in range(100):
+        y = nn.leaky_relu(y)
+    mx.eval(y)
+
+
+def prelu(x: mx.array):
+    y = x
+    for i in range(100):
+        y = nn.prelu(y, mx.ones(1))
+    mx.eval(y)
+
+
+def softplus(x: mx.array):
+    y = x
+    for i in range(100):
+        y = nn.softplus(y)
+    mx.eval(y)
+
+
+def mish(x: mx.array):
+    y = x
+    for i in range(100):
+        y = nn.mish(y)
+    mx.eval(y)
+
+
+def leaky_relu(x):
+    y = x
+    for i in range(100):
+        y = nn.leaky_relu(y)
+    mx.eval(y)
+
+
+def elu(x):
+    y = x
+    for i in range(100):
+        y = nn.elu(y)
+    mx.eval(y)
+
+
+def relu6(x):
+    y = x
+    for i in range(100):
+        y = nn.relu6(y)
+    mx.eval(y)
+
+
+def softplus(x):
+    y = x
+    for i in range(100):
+        y = nn.softplus(y)
+    mx.eval(y)
+
+
+def celu(x):
+    y = x
+    for i in range(100):
+        y = nn.celu(y)
+    mx.eval(y)
+
+
+def log_sigmoid(x):
+    y = x
+    for i in range(100):
+        y = nn.log_sigmoid(y)
     mx.eval(y)
 
 
@@ -180,6 +270,20 @@ def topk(axis, x):
     mx.eval(ys)
 
 
+def step_function(x):
+    y = x
+    for i in range(100):
+        y = nn.step(x)
+    mx.eval(y)
+
+
+def selu(x):
+    y = x
+    for i in range(100):
+        y = nn.selu(x)
+    mx.eval(y)
+
+
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument("benchmark", help="Choose the benchmark to run")
@@ -211,9 +315,7 @@ if __name__ == "__main__":
     parser.add_argument(
         "--fused", action="store_true", help="Use fused functions where possible"
     )
-    parser.add_argument(
-        "--dtype", choices=["float32", "float16", "bfloat16"], default="float32"
-    )
+    parser.add_argument("--dtype", type=dtype_from_str, default=[], action="append")
 
     args = parser.parse_args()
 
@@ -230,11 +332,15 @@ if __name__ == "__main__":
         mx.set_default_device(mx.cpu)
     else:
         mx.set_default_device(mx.gpu)
-    dtype = dict(float32=mx.float32, float16=mx.float16, bfloat16=mx.bfloat16)[
-        args.dtype
-    ]
+
+    types = args.dtype
+    if not types:
+        types = [mx.float32]
+    if len(types) < len(args.size):
+        types = types + [types[0]] * (len(args.size) - len(types))
+
     xs = []
-    for size in args.size:
+    for size, dtype in zip(args.size, types):
         xs.append(mx.random.normal(size).astype(dtype))
     for i, t in enumerate(args.transpose):
         if t is None:
@@ -250,6 +356,9 @@ if __name__ == "__main__":
     elif args.benchmark == "matmul":
         print(bench(matmul, *xs))
 
+    elif args.benchmark == "quant_matmul":
+        print(bench(quant_matmul, *xs))
+
     elif args.benchmark == "linear":
         print(bench(linear, *xs))
 
@@ -277,6 +386,26 @@ if __name__ == "__main__":
     elif args.benchmark == "relu":
         print(bench(relu, x))
 
+    elif args.benchmark == "elu":
+        print(bench(elu, x))
+
+    elif args.benchmark == "relu6":
+        print(bench(relu6, x))
+
+    elif args.benchmark == "celu":
+        print(bench(celu, x))
+
+    elif args.benchmark == "log_sigmoid":
+        print(bench(log_sigmoid, x))
+
+    elif args.benchmark == "leaky_relu":
+        print(bench(leaky_relu, x))
+    elif args.benchmark == "prelu":
+        print(bench(prelu, x))
+    elif args.benchmark == "softplus":
+        print(bench(softplus, x))
+    elif args.benchmark == "mish":
+        print(bench(mish, x))
     elif args.benchmark == "scalar_mul":
         print(bench(scalar_mult, x))
 
@@ -311,5 +440,11 @@ if __name__ == "__main__":
     elif args.benchmark == "topk":
         print(bench(topk, axis, x))
 
+    elif args.benchmark == "step":
+        print(bench(step_function, x))
+
+    elif args.benchmark == "selu":
+        print(bench(selu, x))
+
     else:
         raise ValueError("Unknown benchmark")

benchmarks/python/comparative/bench_torch.py

@@ -22,6 +22,16 @@ def none_or_list(x):
         return [int(xi) for xi in x.split(",")]
 
 
+def dtype_from_str(x):
+    if x == "":
+        return torch.float32
+    else:
+        dt = getattr(torch, x)
+        if not isinstance(dt, torch.dtype):
+            raise ValueError(f"{x} is not a torch dtype")
+        return dt
+
+
 def bench(f, *args):
     for i in range(10):
         f(*args)
@@ -115,6 +125,70 @@ def relu(x):
     sync_if_needed(x)
 
 
+@torch.no_grad()
+def leaky_relu(x):
+    y = x
+    for i in range(100):
+        y = torch.nn.functional.leaky_relu(y)
+    sync_if_needed(x)
+
+
+@torch.no_grad()
+def elu(x):
+    y = x
+    for i in range(100):
+        y = torch.nn.functional.elu(y)
+    sync_if_needed(x)
+
+
+@torch.no_grad()
+def celu(x):
+    y = x
+    for i in range(100):
+        y = torch.nn.functional.celu(y)
+    sync_if_needed(x)
+
+
+@torch.no_grad()
+def relu6(x):
+    y = x
+    for i in range(100):
+        y = torch.nn.functional.relu6(y)
+    sync_if_needed(x)
+
+
+@torch.no_grad()
+def softplus(x):
+    y = x
+    for i in range(100):
+        y = torch.nn.functional.softplus(y)
+    sync_if_needed(x)
+
+
+@torch.no_grad()
+def log_sigmoid(x):
+    y = x
+    for i in range(100):
+        y = torch.nn.functional.logsigmoid(y)
+    sync_if_needed(x)
+
+
+@torch.no_grad()
+def prelu(x: torch.Tensor) -> torch.Tensor:
+    y = x
+    for _ in range(100):
+        y = torch.nn.functional.prelu(y, torch.ones(1).to(y.device))
+    sync_if_needed(x)
+
+
+@torch.no_grad()
+def mish(x: torch.Tensor) -> torch.Tensor:
+    y = x
+    for _ in range(100):
+        return torch.nn.functional.mish(y)
+    sync_if_needed(x)
+
+
 @torch.no_grad()
 def scalar_mult(x):
     y = x
@@ -209,6 +283,14 @@ def topk(axis, x):
     sync_if_needed(x)
 
 
+@torch.no_grad()
+def selu(x):
+    y = x
+    for i in range(100):
+        y = torch.nn.functional.selu(y)
+    sync_if_needed(x)
+
+
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument("benchmark", help="Choose the benchmark to run")
@@ -240,7 +322,7 @@ if __name__ == "__main__":
     parser.add_argument(
         "--fused", action="store_true", help="Use fused functions where possible"
     )
-    parser.add_argument("--dtype", choices=["float32", "float16"], default="float32")
+    parser.add_argument("--dtype", type=dtype_from_str, default=[], action="append")
 
     args = parser.parse_args()
 
@@ -255,9 +337,15 @@ if __name__ == "__main__":
 
     torch.set_num_threads(1)
     device = "cpu" if args.cpu else "mps"
-    dtype = dict(float32=torch.float32, float16=torch.float16)[args.dtype]
+
+    types = args.dtype
+    if not types:
+        types = [torch.float32]
+    if len(types) < len(args.size):
+        types = types + [types[0]] * (len(args.size) - len(types))
+
     xs = []
-    for size in args.size:
+    for size, dtype in zip(args.size, types):
         xs.append(torch.randn(*size).to(device).to(dtype))
     for i, t in enumerate(args.transpose):
         if t is None:
@@ -302,6 +390,28 @@ if __name__ == "__main__":
     elif args.benchmark == "relu":
         print(bench(relu, x))
 
+    elif args.benchmark == "leaky_relu":
+        print(bench(leaky_relu, x))
+
+    elif args.benchmark == "elu":
+        print(bench(elu, x))
+
+    elif args.benchmark == "relu6":
+        print(bench(relu6, x))
+
+    elif args.benchmark == "softplus":
+        print(bench(softplus, x))
+
+    elif args.benchmark == "celu":
+        print(bench(celu, x))
+
+    elif args.benchmark == "log_sigmoid":
+        print(bench(log_sigmoid, x))
+
+    elif args.benchmark == "prelu":
+        print(bench(prelu, x))
+    elif args.benchmark == "mish":
+        print(bench(mish, x))
     elif args.benchmark == "scalar_mul":
         print(bench(scalar_mult, x))
 

benchmarks/python/comparative/compare.py

@@ -193,6 +193,27 @@ if __name__ == "__main__":
     compare_filtered("softmax --size 2x1024x1024 --axis 1 --fused --cpu")
     compare_filtered("relu --size 32x16x1024")
     compare_filtered("relu --size 32x16x1024 --cpu")
+    compare_filtered("leaky_relu --size 32x16x1024")
+    compare_filtered("leaky_relu --size 32x16x1024 --cpu")
+    compare_filtered("elu --size 32x16x1024")
+    compare_filtered("elu --size 32x16x1024 --cpu")
+    compare_filtered("relu6 --size 32x16x1024")
+    compare_filtered("relu6 --size 32x16x1024 --cpu")
+    compare_filtered("softplus --size 32x16x1024")
+    compare_filtered("softplus --size 32x16x1024 --cpu")
+    compare_filtered("celu --size 32x16x1024")
+    compare_filtered("celu --size 32x16x1024 --cpu")
+    compare_filtered("log_sigmoid --size 32x16x1024")
+    compare_filtered("log_sigmoid --size 32x16x1024 --cpu")
+    compare_filtered("step --size 32x16x1024")
+    compare_filtered("step --size 32x16x1024 --cpu")
+    compare_filtered("selu --size 32x16x1024")
+    compare_filtered("selu --size 32x16x1024 --cpu")
+    # compare_filtered("mish --size 32x16x1024") NOTE: Torch does not implement Mish in MPS atm
+    compare_filtered("mish --size 32x16x1024 --cpu")
+    compare_filtered("prelu --size 32x16x1024")
+    compare_filtered("prelu --size 32x16x1024 --cpu")
+
     compare_filtered("scalar_mul --size 32x16x1024")
     compare_filtered("scalar_mul --size 32x16x1024 --cpu")
     compare_filtered("cross_entropy --size 256x1024")

benchmarks/python/llama_torch_bench.py

@@ -4,8 +4,8 @@ import math
 import time
 
 import torch
-import torch.nn as nn
 import torch.mps
+import torch.nn as nn
 
 
 def sync_if_needed(x):

benchmarks/python/single_ops.py

@@ -1,8 +1,8 @@
 # Copyright © 2023 Apple Inc.
 
 import argparse
-import mlx.core as mx
 
+import mlx.core as mx
 from time_utils import time_fn
 
 

docs/.gitignore

@@ -1 +1,2 @@
 src/python/_autosummary*/
+src/python/nn/_autosummary*/

docs/README.md

@@ -7,7 +7,7 @@ for example with `conda`:
 
 ```
 conda install sphinx
-pip install sphinx-rtd-theme
+pip install sphinx-book-theme
 ```
 
 ### Build

docs/src/conf.py

@@ -10,8 +10,8 @@ import subprocess
 project = "MLX"
 copyright = "2023, MLX Contributors"
 author = "MLX Contributors"
-version = "0.0.0"
-release = "0.0.0"
+version = "0.0.6"
+release = "0.0.6"
 
 # -- General configuration ---------------------------------------------------
 
@@ -39,7 +39,17 @@ pygments_style = "sphinx"
 
 # -- Options for HTML output -------------------------------------------------
 
-html_theme = "sphinx_rtd_theme"
+html_theme = "sphinx_book_theme"
+
+html_theme_options = {
+    "show_toc_level": 2,
+    "repository_url": "https://github.com/ml-explore/mlx",
+    "use_repository_button": True,
+    "navigation_with_keys": False,
+}
+
+html_logo = "_static/mlx_logo.png"
+
 
 # -- Options for HTMLHelp output ---------------------------------------------
 

docs/src/dev/extensions.rst

@@ -131,7 +131,7 @@ back and go to our example to give ourselves a more concrete image.
         * A primitive must know how to evaluate itself on the CPU/GPU
         * for the given inputs and populate the output array.
         *
-        * To avoid unecessary allocations, the evaluation function
+        * To avoid unnecessary allocations, the evaluation function
         * is responsible for allocating space for the array.
         */
         void eval_cpu(const std::vector<array>& inputs, array& out) override;
@@ -150,7 +150,7 @@ back and go to our example to give ourselves a more concrete image.
             const std::vector<int>& argnums) override;
 
         /**
-        * The primitive must know how to vectorize itself accross
+        * The primitive must know how to vectorize itself across
         * the given axes. The output is a pair containing the array
         * representing the vectorized computation and the axis which
         * corresponds to the output vectorized dimension.
@@ -945,4 +945,4 @@ Scripts
 .. _Metal-cpp: https://developer.apple.com/metal/cpp/
 .. _`Metal Specification`: https://developer.apple.com/metal/Metal-Shading-Language-Specification.pdf
 .. _`Metal Example`: https://developer.apple.com/documentation/metal/performing_calculations_on_a_gpu?language=objc
-.. _PyBind11: https://pybind11.readthedocs.io/en/stable/
+.. _PyBind11: https://pybind11.readthedocs.io/en/stable/

docs/src/examples/llama-inference.rst

@@ -321,7 +321,7 @@ which can then be used to update the model. Note that the method above incurs
 several unnecessary copies from disk to numpy and then from numpy to MLX. It
 will be replaced in the future with direct loading to MLX.
 
-You can download the full example code in `mlx-examples <code>`_. Assuming, the
+You can download the full example code in `mlx-examples`_. Assuming, the
 existence of ``weights.pth`` and ``tokenizer.model`` in the current working
 directory we can play around with our inference script as follows (the timings
 are representative of an M1 Ultra and the 7B parameter Llama model):
@@ -369,9 +369,9 @@ Scripts
 
 .. admonition:: Download the code
 
-   The full example code is available in `mlx-examples <code>`_.
+   The full example code is available in `mlx-examples`_.
 
-.. code: `https://github.com/ml-explore/mlx-examples/tree/main/llama`_
+.. _mlx-examples: https://github.com/ml-explore/mlx-examples/tree/main/llama
 
 .. [1] Su, J., Lu, Y., Pan, S., Murtadha, A., Wen, B. and Liu, Y., 2021.
    Roformer: Enhanced transformer with rotary position embedding. arXiv

docs/src/examples/mlp.rst

@@ -61,7 +61,10 @@ set:
   def eval_fn(model, X, y):
       return mx.mean(mx.argmax(model(X), axis=1) == y)
 
-Next, setup the problem parameters and load the data:
+Next, setup the problem parameters and load the data. To load the data, you need our
+`mnist data loader
+<https://github.com/ml-explore/mlx-examples/blob/main/mnist/mnist.py>`_, which
+we will import as `mnist`.
 
 .. code-block:: python
 
@@ -127,5 +130,5 @@ Finally, we put it all together by instantiating the model, the
   This should not be confused with :func:`mlx.core.value_and_grad`.
 
 The model should train to a decent accuracy (about 95%) after just a few passes
-over the training set. The `full example <https://github.com/ml-explore/mlx-examples/tree/main/mlp>`_
+over the training set. The `full example <https://github.com/ml-explore/mlx-examples/tree/main/mnist>`_
 is available in the MLX GitHub repo.

docs/src/index.rst

@@ -1,6 +1,30 @@
 MLX
 ===
 
+MLX is a NumPy-like array framework designed for efficient and flexible machine
+learning on Apple silicon, brought to you by Apple machine learning research.
+
+The Python API closely follows NumPy with a few exceptions. MLX also has a
+fully featured C++ API which closely follows the Python API.
+
+The main differences between MLX and NumPy are:
+
+ - **Composable function transformations**: MLX has composable function
+   transformations for automatic differentiation, automatic vectorization,
+   and computation graph optimization.
+ - **Lazy computation**: Computations in MLX are lazy. Arrays are only
+   materialized when needed.
+ - **Multi-device**: Operations can run on any of the supported devices (CPU,
+   GPU, ...)
+
+The design of MLX is inspired by frameworks like `PyTorch
+<https://pytorch.org/>`_, `Jax <https://github.com/google/jax>`_, and
+`ArrayFire <https://arrayfire.org/>`_. A noteable difference from these
+frameworks and MLX is the *unified memory model*. Arrays in MLX live in shared
+memory. Operations on MLX arrays can be performed on any of the supported
+device types without performing data copies. Currently supported device types
+are the CPU and GPU.
+
 .. toctree::
    :caption: Install
    :maxdepth: 1
@@ -12,6 +36,7 @@ MLX
    :maxdepth: 1
 
    quick_start
+   unified_memory
    using_streams
 
 .. toctree::

docs/src/install.rst

@@ -11,6 +11,33 @@ silicon computer is
 
     pip install mlx
 
+To install from PyPI you must meet the following requirements:
+
+- Using an M series chip (Apple silicon)
+- Using a native Python >= 3.8
+- macOS >= 13.3
+
+.. note::
+    MLX is only available on devices running macOS >= 13.3 
+    It is highly recommended to use macOS 14 (Sonoma)
+
+Troubleshooting
+^^^^^^^^^^^^^^^
+
+*My OS and Python versions are in the required range but pip still does not find
+a matching distribution.*
+
+Probably you are using a non-native Python. The output of
+
+.. code-block:: shell
+
+  python -c "import platform; print(platform.processor())"
+
+should be ``arm``. If it is ``i386`` (and you have M series machine) then you
+are using a non-native Python. Switch your Python to a native Python. A good
+way to do this is with `Conda <https://stackoverflow.com/q/65415996>`_.
+
+
 Build from source
 -----------------
 
@@ -19,6 +46,7 @@ Build Requirements
 
 - A C++ compiler with C++17 support (e.g. Clang >= 5.0)
 - `cmake <https://cmake.org/>`_ -- version 3.24 or later, and ``make``
+- Xcode >= 14.3 (Xcode >= 15.0 for macOS 14 and above)
 
 
 Python API
@@ -56,8 +84,16 @@ To make sure the install is working run the tests with:
 
 .. code-block:: shell
 
+  pip install ".[testing]"
   python -m unittest discover python/tests
 
+Optional: Install stubs to enable auto completions and type checking from your IDE:
+
+.. code-block:: shell
+
+  pip install ".[dev]"
+  python setup.py generate_stubs
+
 C++ API
 ^^^^^^^
 
@@ -111,3 +147,45 @@ should point to the path to the built metal library.
      - ON
    * - MLX_BUILD_PYTHON_BINDINGS
      - OFF
+
+
+.. note::
+
+    If you have multiple Xcode installations and wish to use 
+    a specific one while building, you can do so by adding the 
+    following environment variable before building 
+
+    .. code-block:: shell
+
+      export DEVELOPER_DIR="/path/to/Xcode.app/Contents/Developer/"
+
+    Further, you can use the following command to find out which 
+    macOS SDK will be used
+
+    .. code-block:: shell
+
+      xcrun -sdk macosx --show-sdk-version
+
+Troubleshooting
+^^^^^^^^^^^^^^^
+
+Metal not found
+~~~~~~~~~~~~~~~
+
+You see the following error when you try to build:
+
+.. code-block:: shell
+
+  error: unable to find utility "metal", not a developer tool or in PATH
+
+To fix this, first make sure you have Xcode installed:
+
+.. code-block:: shell
+
+  xcode-select --install
+
+Then set the active developer directory:
+
+.. code-block:: shell
+
+  sudo xcode-select --switch /Applications/Xcode.app/Contents/Developer

docs/src/python/array.rst

@@ -34,6 +34,7 @@ Array
     array.prod
     array.reciprocal
     array.reshape
+    array.round
     array.rsqrt
     array.sin
     array.split

docs/src/python/nn.rst

@@ -64,7 +64,6 @@ Quick Start with Neural Networks
     # gradient with respect to `mlp.trainable_parameters()`
     loss_and_grad = nn.value_and_grad(mlp, l2_loss)
 
-
 .. _module_class:
 
 The Module Class
@@ -86,20 +85,58 @@ name should not start with ``_``). It can be arbitrarily nested in other
 :meth:`Module.parameters` can be used to extract a nested dictionary with all
 the parameters of a module and its submodules.
 
-A :class:`Module` can also keep track of "frozen" parameters.
-:meth:`Module.trainable_parameters` returns only the subset of
-:meth:`Module.parameters` that is not frozen. When using
-:meth:`mlx.nn.value_and_grad` the gradients returned will be with respect to these
-trainable parameters.
+A :class:`Module` can also keep track of "frozen" parameters. See the
+:meth:`Module.freeze` method for more details. :meth:`mlx.nn.value_and_grad`
+the gradients returned will be with respect to these trainable parameters.
 
-Updating the parameters
+
+Updating the Parameters
 ^^^^^^^^^^^^^^^^^^^^^^^
 
 MLX modules allow accessing and updating individual parameters. However, most
 times we need to update large subsets of a module's parameters. This action is
-performed by :meth:`Module.update`. 
+performed by :meth:`Module.update`.
 
-Value and grad
+
+Inspecting Modules
+^^^^^^^^^^^^^^^^^^
+
+The simplest way to see the model architecture is to print it. Following along with
+the above example, you can print the ``MLP`` with:
+
+.. code-block:: python
+
+  print(mlp)
+
+This will display:
+
+.. code-block:: shell
+
+  MLP(
+    (layers.0): Linear(input_dims=2, output_dims=128, bias=True)
+    (layers.1): Linear(input_dims=128, output_dims=128, bias=True)
+    (layers.2): Linear(input_dims=128, output_dims=10, bias=True)
+  )
+
+To get more detailed information on the arrays in a :class:`Module` you can use
+:func:`mlx.utils.tree_map` on the parameters. For example, to see the shapes of
+all the parameters in a :class:`Module` do:
+
+.. code-block:: python
+    
+   from mlx.utils import tree_map
+   shapes = tree_map(lambda p: p.shape, mlp.parameters())
+
+As another example, you can count the number of parameters in a :class:`Module`
+with:
+
+.. code-block:: python
+    
+   from mlx.utils import tree_flatten
+   num_params = sum(v.size for _, v in tree_flatten(mlp.parameters()))
+
+
+Value and Grad
 --------------
 
 Using a :class:`Module` does not preclude using MLX's high order function
@@ -133,40 +170,14 @@ In detail:
   :meth:`mlx.core.value_and_grad`
 
 .. autosummary::
+   :recursive:
    :toctree: _autosummary
 
    value_and_grad
+   Module
 
-Neural Network Layers
----------------------
+.. toctree::
 
-.. autosummary::
-   :toctree: _autosummary
-   :template: nn-module-template.rst
-
-   Embedding
-   ReLU
-   GELU
-   SiLU
-   Linear
-   Conv1d
-   Conv2d
-   LayerNorm
-   RMSNorm
-   GroupNorm
-   RoPE
-   MultiHeadAttention
-   Sequential
-
-Layers without parameters (e.g. activation functions) are also provided as
-simple functions.
-
-.. autosummary::
-   :toctree: _autosummary_functions
-   :template: nn-module-template.rst
-
-   gelu
-   gelu_approx
-   gelu_fast_approx
-   relu
-   silu
+   nn/layers
+   nn/functions
+   nn/losses

docs/src/python/nn/functions.rst

@@ -0,0 +1,23 @@
+.. _nn_functions:
+
+.. currentmodule:: mlx.nn
+
+Functions
+---------
+
+Layers without parameters (e.g. activation functions) are also provided as
+simple functions.
+
+.. autosummary::
+   :toctree: _autosummary_functions
+   :template: nn-module-template.rst
+
+   gelu
+   gelu_approx
+   gelu_fast_approx
+   relu
+   prelu
+   silu
+   step
+   selu
+   mish

docs/src/python/nn/layers.rst

@@ -0,0 +1,29 @@
+.. _layers:
+
+.. currentmodule:: mlx.nn
+
+Layers
+------
+
+.. autosummary::
+   :toctree: _autosummary
+   :template: nn-module-template.rst
+
+   Embedding
+   ReLU
+   PReLU
+   GELU
+   SiLU
+   Step
+   SELU
+   Mish
+   Linear
+   Conv1d
+   Conv2d
+   LayerNorm
+   RMSNorm
+   GroupNorm
+   RoPE
+   MultiHeadAttention
+   Sequential
+   QuantizedLinear

docs/src/python/nn/losses.rst

@@ -0,0 +1,19 @@
+.. _losses:
+
+.. currentmodule:: mlx.nn.losses
+
+Loss Functions
+--------------
+
+.. autosummary::
+   :toctree: _autosummary_functions
+   :template: nn-module-template.rst
+
+   binary_cross_entropy
+   cross_entropy
+   kl_div_loss
+   l1_loss
+   mse_loss
+   nll_loss
+   smooth_l1_loss
+   triplet_loss

docs/src/python/nn/module.rst

@@ -1,7 +0,0 @@
-mlx.nn.Module
-=============
-
-.. currentmodule:: mlx.nn
-
-.. autoclass:: Module
-   :members:

docs/src/python/ops.rst

@@ -26,23 +26,32 @@ Operations
    argsort
    array_equal
    broadcast_to
+   ceil
+   clip
    concatenate
    convolve
    conv1d
    conv2d
    cos
    cosh
+   dequantize
    divide
    equal
    erf
    erfinv
    exp
    expand_dims
+   eye
+   flatten
+   floor
+   floor_divide
    full
    greater
    greater_equal
+   identity
    less
    less_equal
+   linspace
    load
    log
    log2
@@ -57,6 +66,7 @@ Operations
    mean
    min
    minimum
+   moveaxis
    multiply
    negative
    ones
@@ -64,8 +74,11 @@ Operations
    partition
    pad
    prod
+   quantize
+   quantized_matmul
    reciprocal
    reshape
+   round
    rsqrt
    save
    savez
@@ -80,14 +93,19 @@ Operations
    sqrt
    square
    squeeze
+   stack
    stop_gradient
    subtract
    sum
+   swapaxes
    take
    take_along_axis
    tan
    tanh
    transpose
+   tri
+   tril
+   triu
    var
    where
    zeros

docs/src/python/optimizers.rst

@@ -38,4 +38,10 @@ model's parameters and the **optimizer state**.
    OptimizerState
    Optimizer
    SGD
+   RMSprop
+   Adagrad
+   AdaDelta
    Adam
+   AdamW
+   Adamax
+   Lion

docs/src/python/transforms.rst

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

docs/src/quick_start.rst

@@ -1,28 +1,6 @@
 Quick Start Guide
 =================
 
-MLX is a NumPy-like array framework designed for efficient and flexible
-machine learning on Apple silicon. The Python API closely follows NumPy with
-a few exceptions. MLX also has a fully featured C++ API which closely follows
-the Python API.
-
-The main differences between MLX and NumPy are: 
-
- - **Composable function transformations**: MLX has composable function
-   transformations for automatic differentiation, automatic vectorization,
-   and computation graph optimization.
- - **Lazy computation**: Computations in MLX are lazy. Arrays are only
-   materialized when needed.
- - **Multi-device**: Operations can run on any of the supported devices (CPU,
-   GPU, ...) 
-
-The design of MLX is strongly inspired by frameworks like `PyTorch
-<https://pytorch.org/>`_, `Jax <https://github.com/google/jax>`_, and
-`ArrayFire <https://arrayfire.org/>`_. A noteable difference from these
-frameworks and MLX is the *unified memory model*. Arrays in MLX live in shared
-memory. Operations on MLX arrays can be performed on any of the supported
-device types without performing data copies. Currently supported device types
-are the CPU and GPU.
 
 Basics
 ------
@@ -84,10 +62,3 @@ and :func:`jvp` for Jacobian-vector products.
 
 Use :func:`value_and_grad` to efficiently compute both a function's output and
 gradient with respect to the function's input. 
-
-
-Devices and Streams 
--------------------
-
-
-

docs/src/unified_memory.rst

@@ -0,0 +1,78 @@
+.. _unified_memory:
+
+Unified Memory
+==============
+
+.. currentmodule:: mlx.core
+
+Apple silicon has a unified memory architecture. The CPU and GPU have direct
+access to the same memory pool. MLX is designed to take advantage of that.
+
+Concretely, when you make an array in MLX you don't have to specify its location:
+
+
+.. code-block:: python
+
+  a = mx.random.normal((100,))
+  b = mx.random.normal((100,))
+
+Both ``a`` and ``b`` live in unified memory.
+
+In MLX, rather than moving arrays to devices, you specify the device when you
+run the operation. Any device can perform any operation on ``a`` and ``b``
+without needing to move them from one memory location to another. For example: 
+
+.. code-block:: python
+
+  mx.add(a, b, stream=mx.cpu)
+  mx.add(a, b, stream=mx.gpu)
+
+In the above, both the CPU and the GPU will perform the same add
+operation. The operations can (and likely will) be run in parallel since
+there are no dependencies between them. See :ref:`using_streams` for more
+information the semantics of streams in MLX.
+
+In the above ``add`` example, there are no dependencies between operations, so
+there is no possibility for race conditions. If there are dependencies, the
+MLX scheduler will automatically manage them. For example:
+
+.. code-block:: python
+
+  c = mx.add(a, b, stream=mx.cpu)
+  d = mx.add(a, c, stream=mx.gpu)
+
+In the above case, the second ``add`` runs on the GPU but it depends on the
+output of the first ``add`` which is running on the CPU. MLX will
+automatically insert a dependency between the two streams so that the second
+``add`` only starts executing after the first is complete and ``c`` is
+available.
+
+A Simple Example
+~~~~~~~~~~~~~~~~
+
+Here is a more interesting (albeit slightly contrived example) of how unified
+memory can be helpful. Suppose we have the following computation:
+
+.. code-block:: python
+
+  def fun(a, b, d1, d2):
+    x = mx.matmul(a, b, stream=d1)
+    for _ in range(500):
+        b = mx.exp(b, stream=d2)
+    return x, b
+
+which we want to run with the following arguments:
+
+.. code-block:: python
+
+  a = mx.random.uniform(shape=(4096, 512))
+  b = mx.random.uniform(shape=(512, 4))
+
+The first ``matmul`` operation is a good fit for the GPU since it's more
+compute dense. The second sequence of operations are a better fit for the CPU,
+since they are very small and would probably be overhead bound on the GPU.
+
+If we time the computation fully on the GPU, we get 2.8 milliseconds. But if we
+run the computation with ``d1=mx.gpu`` and ``d2=mx.cpu``, then the time is only
+about 1.4 milliseconds, about twice as fast. These times were measured on an M1
+Max.

docs/src/using_streams.rst

@@ -1,3 +1,5 @@
+.. _using_streams:
+
 Using Streams
 =============
 

examples/extensions/axpby/axpby.h

@@ -58,7 +58,7 @@ class Axpby : public Primitive {
       const std::vector<int>& argnums) override;
 
   /**
-   * The primitive must know how to vectorize itself accross
+   * The primitive must know how to vectorize itself across
    * the given axes. The output is a pair containing the array
    * representing the vectorized computation and the axis which
    * corresponds to the output vectorized dimension.

examples/extensions/mlx_sample_extensions/__init__.py

@@ -1,4 +1,5 @@
 # Copyright © 2023 Apple Inc.
 
 import mlx.core as mx
+
 from .mlx_sample_extensions import *

examples/extensions/setup.py

@@ -1,8 +1,9 @@
 # Copyright © 2023 Apple Inc.
 
-from mlx import extension
 from setuptools import setup
 
+from mlx import extension
+
 if __name__ == "__main__":
     setup(
         name="mlx_sample_extensions",
@@ -14,5 +15,5 @@ if __name__ == "__main__":
         package_dir={"": "."},
         package_data={"mlx_sample_extensions": ["*.so", "*.dylib", "*.metallib"]},
         zip_safe=False,
-        python_requires=">=3.7",
+        python_requires=">=3.8",
     )

examples/python/linear_regression.py

@@ -1,8 +1,9 @@
 # Copyright © 2023 Apple Inc.
 
-import mlx.core as mx
 import time
 
+import mlx.core as mx
+
 num_features = 100
 num_examples = 1_000
 num_iters = 10_000

examples/python/logistic_regression.py

@@ -1,8 +1,9 @@
 # Copyright © 2023 Apple Inc.
 
-import mlx.core as mx
 import time
 
+import mlx.core as mx
+
 num_features = 100
 num_examples = 1_000
 num_iters = 10_000

mlx/array.h

@@ -154,8 +154,8 @@ class array {
     };
 
    private:
-    int idx;
     const array& arr;
+    int idx;
   };
 
   ArrayIterator begin() const {

mlx/backend/accelerate/CMakeLists.txt

@@ -4,6 +4,7 @@ target_sources(
   ${CMAKE_CURRENT_SOURCE_DIR}/conv.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/matmul.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/primitives.cpp
+  ${CMAKE_CURRENT_SOURCE_DIR}/quantized.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/reduce.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/softmax.cpp
 )

mlx/backend/accelerate/primitives.cpp

@@ -26,12 +26,14 @@ DEFAULT(ArgReduce)
 DEFAULT(ArgSort)
 DEFAULT(AsStrided)
 DEFAULT(Broadcast)
+DEFAULT(Ceil)
 DEFAULT(Concatenate)
 DEFAULT(Copy)
 DEFAULT(Equal)
 DEFAULT(Erf)
 DEFAULT(ErfInv)
 DEFAULT(FFT)
+DEFAULT(Floor)
 DEFAULT(Gather)
 DEFAULT(Greater)
 DEFAULT(GreaterEqual)
@@ -45,6 +47,7 @@ DEFAULT(Pad)
 DEFAULT(Partition)
 DEFAULT(RandomBits)
 DEFAULT(Reshape)
+DEFAULT(Round)
 DEFAULT(Scatter)
 DEFAULT(Sigmoid)
 DEFAULT(Sign)
@@ -322,6 +325,45 @@ void Divide::eval_cpu(const std::vector<array>& inputs, array& out) {
   }
 }
 
+// TODO: Avoid code duplication with the common backend.
+struct RemainderFn {
+  template <typename T>
+  std::enable_if_t<!std::is_integral_v<T>, T> operator()(
+      T numerator,
+      T denominator) {
+    return std::fmod(numerator, denominator);
+  }
+
+  template <typename T>
+  std::enable_if_t<std::is_integral_v<T>, T> operator()(
+      T numerator,
+      T denominator) {
+    return numerator % denominator;
+  }
+};
+
+void Remainder::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 2);
+  auto& a = inputs[0];
+  auto& b = inputs[1];
+
+  if (a.dtype() == float32) {
+    binary(
+        a,
+        b,
+        out,
+        RemainderFn{},
+        UseDefaultBinaryOp(),
+        UseDefaultBinaryOp(),
+        [](const auto* a, const auto* b, auto* o, auto n) {
+          int num_el = n;
+          vvremainderf((float*)o, (const float*)a, (const float*)b, &num_el);
+        });
+  } else {
+    binary(a, b, out, RemainderFn{});
+  }
+}
+
 void Exp::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   const auto& in = inputs[0];
@@ -494,7 +536,7 @@ void Power::eval_cpu(const std::vector<array>& inputs, array& out) {
       b.flags().row_contiguous) {
     int size = a.size();
     out.set_data(allocator::malloc_or_wait(out.nbytes()));
-    vvpowf(out.data<float>(), a.data<float>(), b.data<float>(), &size);
+    vvpowf(out.data<float>(), b.data<float>(), a.data<float>(), &size);
   } else {
     eval(inputs, out);
   }

mlx/backend/accelerate/quantized.cpp

@@ -0,0 +1,107 @@
+// Copyright © 2023 Apple Inc.
+
+#include <cassert>
+
+#include <simd/vector.h>
+
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+
+namespace {
+
+void _qmm_t_4_64(
+    float* result,
+    const float* x,
+    const uint32_t* w,
+    const float* scales,
+    const float* biases,
+    int M,
+    int N,
+    int K) {
+  constexpr int bits = 4;
+  constexpr int group_size = 64;
+  constexpr int bitmask = (1 << bits) - 1;
+  constexpr int pack_factor = 32 / bits;
+  constexpr int packs_in_group = group_size / pack_factor;
+  const int Kg = K / group_size;
+  const int Kw = K / pack_factor;
+
+  for (int m = 0; m < M; m++) {
+    const uint32_t* w_local = w;
+    const float* scales_local = scales;
+    const float* biases_local = biases;
+
+    for (int n = 0; n < N; n++) {
+      const simd_float16* x_local = (simd_float16*)x;
+      simd_float16 sum = 0;
+      for (int k = 0; k < K; k += group_size) {
+        float scale = *scales_local++;
+        float bias = *biases_local++;
+
+        for (int kw = 0; kw < packs_in_group; kw += 2) {
+          // TODO: vectorize this properly
+          simd_uint16 wi;
+          for (int e = 0; e < 2; e++) {
+            uint32_t wii = *w_local++;
+            for (int p = 0; p < 8; p++) {
+              wi[e * 8 + p] = wii & bitmask;
+              wii >>= bits;
+            }
+          }
+          simd_float16 wf = simd_float(wi);
+          wf *= scale;
+          wf += bias;
+
+          sum += (*x_local) * wf;
+          x_local++;
+        }
+      }
+
+      *result = simd_reduce_add(sum);
+      result++;
+    }
+
+    x += K;
+  }
+}
+
+} // namespace
+
+void QuantizedMatmul::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 4);
+
+  auto& x = inputs[0];
+  auto& w = inputs[1];
+  auto& scales = inputs[2];
+  auto& biases = inputs[3];
+
+  if (w.strides()[0] != 1) {
+    throw std::runtime_error("The quantized weight should be transposed");
+  }
+
+  if (!x.flags().row_contiguous || !scales.flags().row_contiguous ||
+      !biases.flags().row_contiguous) {
+    throw std::runtime_error("x, scales and biases should be row contiguous.");
+  }
+
+  if (x.dtype() == float32 && bits_ == 4 && group_size_ == 64) {
+    out.set_data(allocator::malloc_or_wait(out.nbytes()));
+    int K = x.shape(-1);
+    int M = x.size() / K;
+    int N = w.shape(1);
+    _qmm_t_4_64(
+        out.data<float>(),
+        x.data<float>(),
+        w.data<uint32_t>(),
+        scales.data<float>(),
+        biases.data<float>(),
+        M,
+        N,
+        K);
+  } else {
+    eval(inputs, out);
+  }
+}
+
+} // namespace mlx::core

mlx/backend/common/CMakeLists.txt

@@ -8,6 +8,7 @@ target_sources(
   ${CMAKE_CURRENT_SOURCE_DIR}/erf.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/fft.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/primitives.cpp
+  ${CMAKE_CURRENT_SOURCE_DIR}/quantized.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/reduce.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/scan.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/softmax.cpp

mlx/backend/common/binary.cpp

@@ -82,6 +82,29 @@ void Divide::eval(const std::vector<array>& inputs, array& out) {
   binary(a, b, out, [](auto x, auto y) { return x / y; });
 }
 
+struct RemainderFn {
+  template <typename T>
+  std::enable_if_t<!std::is_integral_v<T>, T> operator()(
+      T numerator,
+      T denominator) {
+    return std::fmod(numerator, denominator);
+  }
+
+  template <typename T>
+  std::enable_if_t<std::is_integral_v<T>, T> operator()(
+      T numerator,
+      T denominator) {
+    return numerator % denominator;
+  }
+};
+
+void Remainder::eval(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 2);
+  auto& a = inputs[0];
+  auto& b = inputs[1];
+  binary(a, b, out, RemainderFn{});
+}
+
 void Equal::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 2);
   if (equal_nan_) {

mlx/backend/common/conv.cpp

@@ -357,7 +357,7 @@ void explicit_gemm_conv_1D_cpu(
     gemm_out.set_data(allocator::malloc_or_wait(gemm_out.nbytes()));
   }
 
-  // Peform gemm
+  // Perform gemm
   cblas_sgemm(
       CblasRowMajor,
       CblasNoTrans, // no trans A
@@ -459,7 +459,7 @@ void explicit_gemm_conv_2D_cpu(
     gemm_out.set_data(allocator::malloc_or_wait(gemm_out.nbytes()));
   }
 
-  // Peform gemm
+  // Perform gemm
   cblas_sgemm(
       CblasRowMajor,
       CblasNoTrans, // no trans A

mlx/backend/common/default_primitives.cpp

@@ -1,6 +1,10 @@
 // Copyright © 2023 Apple Inc.
 
+#ifdef ACCELERATE_NEW_LAPACK
+#include <vecLib/cblas_new.h>
+#else
 #include <cblas.h>
+#endif
 
 #include "mlx/array.h"
 #include "mlx/backend/common/copy.h"
@@ -29,17 +33,20 @@ DEFAULT(ArgSort)
 DEFAULT(AsType)
 DEFAULT(AsStrided)
 DEFAULT(Broadcast)
+DEFAULT(Ceil)
 DEFAULT(Concatenate)
 DEFAULT(Convolution)
 DEFAULT(Copy)
 DEFAULT(Cos)
 DEFAULT(Cosh)
 DEFAULT(Divide)
+DEFAULT(Remainder)
 DEFAULT(Equal)
 DEFAULT(Erf)
 DEFAULT(ErfInv)
 DEFAULT(Exp)
 DEFAULT(FFT)
+DEFAULT(Floor)
 DEFAULT(Full)
 DEFAULT(Gather)
 DEFAULT(Greater)
@@ -59,9 +66,11 @@ DEFAULT(NotEqual)
 DEFAULT(Pad)
 DEFAULT(Partition)
 DEFAULT(Power)
+DEFAULT(QuantizedMatmul)
 DEFAULT(RandomBits)
 DEFAULT(Reduce)
 DEFAULT(Reshape)
+DEFAULT(Round)
 DEFAULT(Scan)
 DEFAULT(Scatter)
 DEFAULT(Sigmoid)

mlx/backend/common/primitives.cpp

@@ -167,6 +167,17 @@ void Broadcast::eval(const std::vector<array>& inputs, array& out) {
   out.copy_shared_buffer(in, strides, flags, in.data_size());
 }
 
+void Ceil::eval(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  auto& in = inputs[0];
+  if (not is_integral(in.dtype())) {
+    unary_fp(in, out, [](auto x) { return std::ceil(x); });
+  } else {
+    // No-op integer types
+    out.copy_shared_buffer(in);
+  }
+}
+
 void Concatenate::eval(const std::vector<array>& inputs, array& out) {
   std::vector<int> sizes;
   sizes.push_back(0);
@@ -287,6 +298,17 @@ void Exp::eval(const std::vector<array>& inputs, array& out) {
   }
 }
 
+void Floor::eval(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  auto& in = inputs[0];
+  if (not is_integral(in.dtype())) {
+    unary_fp(in, out, [](auto x) { return std::floor(x); });
+  } else {
+    // No-op integer types
+    out.copy_shared_buffer(in);
+  }
+}
+
 void Full::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   auto& in = inputs[0];
@@ -444,6 +466,17 @@ void Reshape::eval(const std::vector<array>& inputs, array& out) {
   }
 }
 
+void Round::eval(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  auto& in = inputs[0];
+  if (not is_integral(in.dtype())) {
+    unary_fp(in, out, RoundOp());
+  } else {
+    // No-op integer types
+    out.copy_shared_buffer(in);
+  }
+}
+
 void Sigmoid::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   const auto& in = inputs[0];

mlx/backend/common/quantized.cpp

@@ -0,0 +1,184 @@
+// Copyright © 2023 Apple Inc.
+
+#include <cassert>
+
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+
+namespace {
+
+template <typename T, int bits, int group_size>
+void _qmm_t(
+    T* result,
+    const T* x,
+    const uint32_t* w,
+    const T* scales,
+    const T* biases,
+    int M,
+    int N,
+    int K) {
+  constexpr int bitmask = (1 << bits) - 1;
+  constexpr int pack_factor = 32 / bits;
+  constexpr int packs_in_group = group_size / pack_factor;
+  const int Kg = K / group_size;
+  const int Kw = K / pack_factor;
+
+  for (int m = 0; m < M; m++) {
+    const uint32_t* w_local = w;
+    const T* scales_local = scales;
+    const T* biases_local = biases;
+
+    for (int n = 0; n < N; n++) {
+      const T* x_local = x;
+      T sum = 0;
+      for (int k = 0; k < K; k += group_size) {
+        T scale = *scales_local++;
+        T bias = *biases_local++;
+
+        for (int kw = 0; kw < packs_in_group; kw++) {
+          uint32_t wi = *w_local++;
+
+#pragma clang loop unroll(full)
+          for (int p = 0; p < pack_factor; p++) {
+            sum += (*x_local++) * (scale * static_cast<T>(wi & bitmask) + bias);
+            wi >>= bits;
+          }
+        }
+      }
+      *result = sum;
+      result++;
+    }
+
+    x += K;
+  }
+}
+
+template <typename T>
+void _qmm_t_dispatch_typed(
+    T* result,
+    const T* x,
+    const uint32_t* w,
+    const T* scales,
+    const T* biases,
+    int M,
+    int N,
+    int K,
+    int group_size,
+    int bits) {
+  switch (bits) {
+    case 2: {
+      switch (group_size) {
+        case 64:
+          return _qmm_t<T, 2, 64>(result, x, w, scales, biases, M, N, K);
+        case 128:
+          return _qmm_t<T, 2, 128>(result, x, w, scales, biases, M, N, K);
+      }
+    }
+    case 4: {
+      switch (group_size) {
+        case 64:
+          return _qmm_t<T, 4, 64>(result, x, w, scales, biases, M, N, K);
+        case 128:
+          return _qmm_t<T, 4, 128>(result, x, w, scales, biases, M, N, K);
+      }
+    }
+    case 8: {
+      switch (group_size) {
+        case 64:
+          return _qmm_t<T, 8, 64>(result, x, w, scales, biases, M, N, K);
+        case 128:
+          return _qmm_t<T, 8, 128>(result, x, w, scales, biases, M, N, K);
+      }
+    }
+  }
+  std::ostringstream msg;
+  msg << "Quantization type not supported. Provided bits=" << bits
+      << " and group_size=" << group_size
+      << ". The supported options are bits in "
+      << "{2, 4, 8} and group_size in {64, 128}.";
+  throw std::invalid_argument(msg.str());
+}
+
+void _qmm_t_dispatch(
+    array out,
+    const array& x,
+    const array& w,
+    const array& scales,
+    const array& biases,
+    int bits,
+    int group_size) {
+  int K = x.shape(-1);
+  int M = x.size() / K;
+  int N = w.shape(1);
+
+  switch (x.dtype()) {
+    case float32:
+      _qmm_t_dispatch_typed<float>(
+          out.data<float>(),
+          x.data<float>(),
+          w.data<uint32_t>(),
+          scales.data<float>(),
+          biases.data<float>(),
+          M,
+          N,
+          K,
+          bits,
+          group_size);
+      break;
+    case float16:
+      _qmm_t_dispatch_typed<float16_t>(
+          out.data<float16_t>(),
+          x.data<float16_t>(),
+          w.data<uint32_t>(),
+          scales.data<float16_t>(),
+          biases.data<float16_t>(),
+          M,
+          N,
+          K,
+          bits,
+          group_size);
+      break;
+    case bfloat16:
+      _qmm_t_dispatch_typed<bfloat16_t>(
+          out.data<bfloat16_t>(),
+          x.data<bfloat16_t>(),
+          w.data<uint32_t>(),
+          scales.data<bfloat16_t>(),
+          biases.data<bfloat16_t>(),
+          M,
+          N,
+          K,
+          bits,
+          group_size);
+      break;
+    default:
+      throw std::invalid_argument(
+          "[quantized_matmul] only floating types are supported");
+  }
+}
+
+} // namespace
+
+void QuantizedMatmul::eval(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 4);
+
+  auto& x = inputs[0];
+  auto& w = inputs[1];
+  auto& scales = inputs[2];
+  auto& biases = inputs[3];
+
+  if (w.strides()[0] != 1) {
+    throw std::runtime_error("The quantized weight should be transposed");
+  }
+
+  if (!x.flags().row_contiguous || !scales.flags().row_contiguous ||
+      !biases.flags().row_contiguous) {
+    throw std::runtime_error("x, scales and biases should be row contiguous.");
+  }
+
+  out.set_data(allocator::malloc_or_wait(out.nbytes()));
+  _qmm_t_dispatch(out, x, w, scales, biases, group_size_, bits_);
+}
+
+} // namespace mlx::core

mlx/backend/common/unary.h

@@ -53,6 +53,17 @@ struct SignOp {
   }
 };
 
+struct RoundOp {
+  template <typename T>
+  T operator()(T x) {
+    return std::round(x);
+  }
+
+  complex64_t operator()(complex64_t x) {
+    return {std::round(x.real()), std::round(x.imag())};
+  }
+};
+
 template <typename T, typename Op>
 void unary_op(const array& a, array& out, Op op) {
   const T* a_ptr = a.data<T>();

mlx/backend/metal/CMakeLists.txt

@@ -10,6 +10,7 @@ target_sources(
   ${CMAKE_CURRENT_SOURCE_DIR}/matmul.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/metal.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/primitives.cpp
+  ${CMAKE_CURRENT_SOURCE_DIR}/quantized.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/scan.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/softmax.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/sort.cpp

mlx/backend/metal/device.cpp

@@ -44,6 +44,25 @@ std::pair<MTL::Library*, NS::Error*> load_library_from_path(
   return std::make_pair(lib, error);
 }
 
+#ifdef SWIFTPM_BUNDLE
+MTL::Library* try_load_bundle(MTL::Device* device, NS::URL* url) {
+  std::string bundle_path = std::string(url->fileSystemRepresentation()) + "/" +
+      SWIFTPM_BUNDLE + ".bundle";
+  auto bundle = NS::Bundle::alloc()->init(
+      NS::String::string(bundle_path.c_str(), NS::UTF8StringEncoding));
+  if (bundle != nullptr) {
+    std::string resource_path =
+        std::string(bundle->resourceURL()->fileSystemRepresentation()) + "/" +
+        "default.metallib";
+    auto [lib, error] = load_library_from_path(device, resource_path.c_str());
+    if (lib) {
+      return lib;
+    }
+  }
+  return nullptr;
+}
+#endif
+
 MTL::Library* load_library(
     MTL::Device* device,
     const std::string& lib_name = "mlx",
@@ -57,6 +76,26 @@ MTL::Library* load_library(
     }
   }
 
+#ifdef SWIFTPM_BUNDLE
+  // try to load from a swiftpm resource bundle -- scan the available bundles to
+  // find one that contains the named bundle
+  {
+    MTL::Library* library =
+        try_load_bundle(device, NS::Bundle::mainBundle()->bundleURL());
+    if (library != nullptr) {
+      return library;
+    }
+    auto bundles = NS::Bundle::allBundles();
+    for (int i = 0, c = (int)bundles->count(); i < c; i++) {
+      auto bundle = reinterpret_cast<NS::Bundle*>(bundles->object(i));
+      library = try_load_bundle(device, bundle->resourceURL());
+      if (library != nullptr) {
+        return library;
+      }
+    }
+  }
+#endif
+
   // Couldn't find it so let's load it from default_mtllib_path
   {
     auto [lib, error] = load_library_from_path(device, lib_path);
@@ -88,6 +127,12 @@ Device::~Device() {
   for (auto& l : library_map_) {
     l.second->release();
   }
+  for (auto& b : buffer_map_) {
+    b.second.second->release();
+  }
+  for (auto& e : encoder_map_) {
+    e.second->release();
+  }
   device_->release();
   pool_->release();
 }

mlx/backend/metal/kernels/CMakeLists.txt

@@ -18,6 +18,7 @@ set(
   "copy"
   "gemm"
   "gemv"
+  "quantized"
   "random"
   "reduce"
   "scan"

mlx/backend/metal/kernels/binary.metal

@@ -14,6 +14,13 @@ struct Divide {
   template <typename T> T operator()(T x, T y) { return x / y; }
 };
 
+struct Remainder {
+  template <typename T> T operator()(T x, T y) { return x % y; }
+  template <> float operator()(float x, float y) { return fmod(x, y); }
+  template <> half operator()(half x, half y) { return fmod(x, y); }
+  template <> bfloat16_t operator()(bfloat16_t x, bfloat16_t y) { return fmod(x, y); }
+};
+
 struct Equal {
   template <typename T> bool operator()(T x, T y) { return x == y; }
 };
@@ -350,7 +357,7 @@ template <typename T, typename U, typename Op>
   instantiate_binary_all(name, complex64, complex64_t, bool, op)
 
 instantiate_binary_types(add, Add)
-instantiate_binary_float(div, Divide)
+instantiate_binary_types(div, Divide)
 instantiate_binary_types_bool(eq, Equal)
 instantiate_binary_types_bool(ge, Greater)
 instantiate_binary_types_bool(geq, GreaterEqual)
@@ -363,6 +370,7 @@ instantiate_binary_types(min, Minimum)
 instantiate_binary_types(mul, Multiply)
 instantiate_binary_types(sub, Subtract)
 instantiate_binary_types(pow, Power)
+instantiate_binary_types(rem, Remainder)
 
 // NaNEqual only needed for floating point types with boolean output
 instantiate_binary_all(naneq, float16, half, bool, NaNEqual)

mlx/backend/metal/kernels/complex.h

@@ -110,3 +110,16 @@ constexpr complex64_t operator-(complex64_t a, complex64_t b) {
 constexpr complex64_t operator*(complex64_t a, complex64_t b) {
   return {a.real * b.real - a.imag * b.imag, a.real * b.imag + a.imag * b.real};
 }
+
+constexpr complex64_t operator/(complex64_t a, complex64_t b) {
+  auto denom = b.real * b.real + b.imag * b.imag;
+  auto x = a.real * b.real + a.imag * b.imag;
+  auto y = a.imag * b.real - a.real * b.imag;
+  return {x / denom, y / denom};
+}
+
+constexpr complex64_t operator%(complex64_t a, complex64_t b) {
+  auto real = a.real - (b.real * static_cast<int64_t>(a.real / b.real));
+  auto imag = a.imag - (b.imag * static_cast<int64_t>(a.imag / b.imag));
+  return {real, imag};
+}

mlx/backend/metal/kernels/gemv.metal

@@ -3,8 +3,9 @@
 #include <metal_stdlib>
 #include <metal_simdgroup>
 
-#include "mlx/backend/metal/kernels/defines.h"
 #include "mlx/backend/metal/kernels/bf16.h"
+#include "mlx/backend/metal/kernels/defines.h"
+#include "mlx/backend/metal/kernels/utils.h"
 
 using namespace metal;
 
@@ -12,52 +13,57 @@ using namespace metal;
 /// Matrix vector multiplication
 ///////////////////////////////////////////////////////////////////////////////
 
-static constant constexpr int SIMD_SIZE = 32;
+#define MLX_MTL_CONST static constant constexpr const
 
-template <typename T, 
-          const int BM, /* Threadgroup rows (in threads) */
-          const int BN, /* Threadgroup cols (in threads) */
-          const int TM, /* Thread rows (in elements) */
-          const int TN> /* Thread cols (in elements) */
-[[kernel]] void gemv(
-    const device T* mat [[buffer(0)]],
-    const device T* in_vec [[buffer(1)]],
-    device T* out_vec [[buffer(2)]], 
-    const constant int& in_vec_size [[buffer(3)]],
-    const constant int& out_vec_size [[buffer(4)]],
-    const constant int& vector_batch_stride [[buffer(5)]],
-    const constant int& matrix_batch_stride [[buffer(6)]],
-    uint3 tid [[threadgroup_position_in_grid]],
-    uint simd_gid [[simdgroup_index_in_threadgroup]],
-    uint simd_lid [[thread_index_in_simdgroup]]) {
+MLX_MTL_CONST int SIMD_SIZE = 32;
 
-    static_assert(BN == SIMD_SIZE, "gemv block must have a width of SIMD_SIZE");
+template <
+  typename T, 
+  const int BM, /* Threadgroup rows (in threads) */
+  const int BN, /* Threadgroup cols (in threads) */
+  const int TM, /* Thread rows (in elements) */
+  const int TN > /* Thread cols (in elements) */ 
+struct GEMVKernel {
 
-    // - The matrix of size (M = out_vec_size, N = in_vec_size) is divided up 
-    //   into blocks of (BM * TM, BN * TN) divided amoung threadgroups
-    // - Every thread works on a block of (TM, TN)
-    // - We assume each thead group is launched with (BN, BM, 1) threads
-    //
-    // 1. A thread loads TN elements each from mat along TM contiguous rows 
-    //      and the corresponding scalar from the vector
-    // 2. The thread then multiplies and adds to accumulate its local result for the block
-    // 3. At the end, each thread has accumulated results over all blocks across the rows
-    //      These are then summed up across the threadgroup
-    // 4. Each threadgroup writes its accumulated BN * TN outputs
-    //
-    // Edge case handling:
-    // - The threadgroup with the largest tid will have blocks that exceed the matrix
-    //   * The blocks that start outside the matrix are never read (thread results remain zero)
-    //   * The last thread that partialy overlaps with the matrix is shifted inwards 
-    //     such that the thread block fits exactly in the matrix
+  static_assert(BN == SIMD_SIZE, "gemv block must have a width of SIMD_SIZE");
+
+  // - The matrix of size (M = out_vec_size, N = in_vec_size) is divided up 
+  //   into blocks of (BM * TM, BN * TN) divided amoung threadgroups
+  // - Every thread works on a block of (TM, TN)
+  // - We assume each thead group is launched with (BN, BM, 1) threads
+  //
+  // 1. A thread loads TN elements each from mat along TM contiguous rows 
+  //      and the corresponding scalar from the vector
+  // 2. The thread then multiplies and adds to accumulate its local result for the block
+  // 3. At the end, each thread has accumulated results over all blocks across the rows
+  //      These are then summed up across the threadgroup
+  // 4. Each threadgroup writes its accumulated BN * TN outputs
+  //
+  // Edge case handling:
+  // - The threadgroup with the largest tid will have blocks that exceed the matrix
+  //   * The blocks that start outside the matrix are never read (thread results remain zero)
+  //   * The last thread that partialy overlaps with the matrix is shifted inwards 
+  //     such that the thread block fits exactly in the matrix
+
+  MLX_MTL_CONST short tgp_mem_size = BN * TN * 2;
+
+  static METAL_FUNC void run(
+      const device T* mat,
+      const device T* in_vec,
+      device T* out_vec, 
+      const constant int& in_vec_size [[buffer(3)]],
+      const constant int& out_vec_size [[buffer(4)]],
+      threadgroup T* tgp_memory [[threadgroup(0)]],
+      uint3 tid [[threadgroup_position_in_grid]],
+      uint3 lid [[thread_position_in_threadgroup]],
+      uint simd_gid [[simdgroup_index_in_threadgroup]],
+      uint simd_lid [[thread_index_in_simdgroup]]) {
+
+    // Appease compiler 
+    (void)lid;
 
-    // Update batch offsets
-    in_vec += tid.z * vector_batch_stride;
-    mat += tid.z * matrix_batch_stride;
-    out_vec += tid.z * out_vec_size;
-    
     // Threadgroup in_vec cache
-    threadgroup T in_vec_block[BN][TN * 2];
+    threadgroup T* in_vec_block = tgp_memory + simd_lid * TN * 2;
 
     // Thread local accumulation results 
     thread T result[TM] = {0};
@@ -69,7 +75,7 @@ template <typename T,
 
     // Exit simdgroup if rows out of bound
     if(out_row >= out_vec_size) 
-        return;
+      return;
 
     // Adjust tail simdgroup to ensure in bound reads
     out_row = out_row + TM <= out_vec_size ? out_row : out_vec_size - TM;
@@ -79,62 +85,304 @@ template <typename T,
 
     // Loop over in_vec in blocks of BN * TN
     for(int bn = simd_lid * TN; bn < in_vec_size; bn += BN * TN) {
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-        // Prefetch in_vector for threadgroup use
-        if(simd_gid == 0) {
-            // Main load loop
-            if(bn + TN <= in_vec_size) {
-                #pragma clang loop unroll(full)
-                for(int tn = 0; tn < TN; tn++) {
-                    in_vec_block[simd_lid][tn] = in_vec[bn + tn];
-                }
-            } else { // Edgecase
-                #pragma clang loop unroll(full)
-                for(int tn = 0; tn < TN; tn++) {
-                    in_vec_block[simd_lid][tn] = bn + tn < in_vec_size ? in_vec[bn + tn] : 0;
-                }
-            }
-        }
-        threadgroup_barrier(mem_flags::mem_threadgroup);
 
-        // Load for all rows
-        #pragma clang loop unroll(full)
+      threadgroup_barrier(mem_flags::mem_threadgroup);
+
+      // Prefetch in_vector for threadgroup use
+      if(simd_gid == 0) {
+        // Main load loop
+        if(bn + TN <= in_vec_size) {
+
+          #pragma clang loop unroll(full)
+          for(int tn = 0; tn < TN; tn++) {
+            in_vec_block[tn] = in_vec[bn + tn];
+          }
+
+        } else { // Edgecase
+
+          #pragma clang loop unroll(full)
+          for(int tn = 0; tn < TN; tn++) {
+            in_vec_block[tn] = bn + tn < in_vec_size ? in_vec[bn + tn] : 0;
+          }
+
+        }
+      }
+
+      threadgroup_barrier(mem_flags::mem_threadgroup);
+
+      // Load for all rows
+      #pragma clang loop unroll(full)
+      for(int tn = 0; tn < TN; tn++) {
+        v_coeff[tn] = in_vec_block[tn];
+      }
+
+      // Per thread work loop
+      #pragma clang loop unroll(full)
+      for(int tm = 0; tm < TM; tm++) {
+
+        // Load for the row 
         for(int tn = 0; tn < TN; tn++) {
-            v_coeff[tn] = in_vec_block[simd_lid][tn];
+          inter[tn] = mat[tm * in_vec_size + bn + tn];
         }
 
-        // Per thread work loop
-        #pragma clang loop unroll(full)
-        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];
-            }
-
-            // Accumulate results
-            for(int tn = 0; tn < TN; tn++) {
-                result[tm] += inter[tn] * v_coeff[tn];
-            }
+        // Accumulate results
+        for(int tn = 0; tn < TN; tn++) {
+          result[tm] += inter[tn] * v_coeff[tn];
         }
+
+      }
     }
 
     // Simdgroup accumulations
     #pragma clang loop unroll(full)
     for(int tm = 0; tm < TM; tm++) {
-        result[tm] = simd_sum(result[tm]);
+      result[tm] = simd_sum(result[tm]);
     }
 
     // Write outputs
     if(simd_lid == 0) {
-        #pragma clang loop unroll(full)
-        for(int tm = 0; tm < TM; tm++) {
-            out_vec[out_row + tm] = result[tm];
-        }
+
+      #pragma clang loop unroll(full)
+      for(int tm = 0; tm < TM; tm++) {
+        out_vec[out_row + tm] = result[tm];
+      }
+
     }
 
+  }
+
+};
+
+///////////////////////////////////////////////////////////////////////////////
+/// Vector matrix multiplication
+///////////////////////////////////////////////////////////////////////////////
+
+template <
+  typename T, 
+  const int BM, /* Threadgroup rows (in threads) */
+  const int BN, /* Threadgroup cols (in threads) */
+  const int TM, /* Thread rows (in elements) */
+  const int TN > /* Thread cols (in elements) */ 
+struct GEMVTKernel {
+
+  // - The matrix of size (M = in_vec_size, N = out_vec_size) is divided up 
+  //   into blocks of (BM * TM, BN * TN) divided amoung threadgroups
+  // - Every thread works on a block of (TM, TN)
+  // - We assume each thead group is launched with (BN, BM, 1) threads
+  //
+  // 1. A thread loads TN elements each from mat along TM contiguous rows 
+  //      and the corresponding scalar from the vector
+  // 2. The thread then multiplies and adds to accumulate its local result for the block
+  // 3. At the end, each thread has accumulated results over all blocks across the rows
+  //      These are then summed up across the threadgroup
+  // 4. Each threadgroup writes its accumulated BN * TN outputs
+  //
+  // Edge case handling:
+  // - The threadgroup with the largest tid will have blocks that exceed the matrix
+  //   * The blocks that start outside the matrix are never read (thread results remain zero)
+  //   * The last thread that partialy overlaps with the matrix is shifted inwards 
+  //     such that the thread block fits exactly in the matrix
+
+
+  MLX_MTL_CONST short tgp_mem_size = BN * BM * TN;
+
+  static METAL_FUNC void run(
+      const device T* mat,
+      const device T* in_vec,
+      device T* out_vec, 
+      const constant int& in_vec_size [[buffer(3)]],
+      const constant int& out_vec_size [[buffer(4)]],
+      threadgroup T* tgp_memory [[threadgroup(0)]],
+      uint3 tid [[threadgroup_position_in_grid]],
+      uint3 lid [[thread_position_in_threadgroup]],
+      uint simd_gid [[simdgroup_index_in_threadgroup]],
+      uint simd_lid [[thread_index_in_simdgroup]]) {
+
+    // Appease compiler 
+    (void)simd_gid;
+    (void)simd_lid;
+
+    // Thread local accumulation results
+    T result[TN] = {0};
+    T inter[TN];
+    T v_coeff[TM];
+
+    // Threadgroup accumulation results
+    threadgroup T* tgp_results = tgp_memory + lid.x * BM * TN;
+
+    int out_col = (tid.x * BN + lid.x) * TN;
+    int in_row = lid.y * TM;
+
+    // Edgecase handling
+    if (out_col < out_vec_size) {
+
+      out_col = out_col + TN < out_vec_size ? out_col : out_vec_size - TN;
+
+      // Per thread accumulation main loop
+      int bm = in_row;
+      for(; bm < in_vec_size; bm += BM * TM) {
+        // Adding a threadgroup_barrier improves performance slightly
+        // This is possibly it may help exploit cache better
+        threadgroup_barrier(mem_flags::mem_none);
+
+        if(bm + TM <= in_vec_size) {
+
+          #pragma clang loop unroll(full)
+          for(int tm = 0; tm < TM; tm++) {
+            v_coeff[tm] = in_vec[bm + tm];
+          }
+
+          #pragma clang loop unroll(full)
+          for(int tm = 0; tm < TM; tm++) {
+            for(int tn = 0; tn < TN; tn++) {
+              inter[tn] = mat[(bm + tm) * out_vec_size + out_col + tn];
+            }
+            for(int tn = 0; tn < TN; tn++) {
+              result[tn] += v_coeff[tm] * inter[tn];
+            }
+          }
+        
+        } else { // Edgecase handling
+          for(int tm = 0; bm + tm < in_vec_size; tm++) {
+            v_coeff[tm] = in_vec[bm + tm];
+
+            for(int tn = 0; tn < TN; tn++) {
+              inter[tn] = mat[(bm + tm) * out_vec_size + out_col + tn];
+            }
+            for(int tn = 0; tn < TN; tn++) {
+              result[tn] += v_coeff[tm] * inter[tn];
+            }
+
+          }
+        }
+      }
+
+    }
+
+    // Threadgroup collection
+
+    #pragma clang loop unroll(full)
+    for(int i = 0; i < TN; i++) {
+      tgp_results[lid.y * TN + i] = result[i];
+    }
+
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    // Threadgroup accumulation and writing out results
+    if(lid.y == 0 && out_col < out_vec_size) {
+      
+      #pragma clang loop unroll(full)
+      for(int i = 1; i < BM; i++) {
+
+        #pragma clang loop unroll(full)
+        for(int j = 0; j < TN; j++) {
+          result[j] += tgp_results[i * TN + j];
+        }
+      }
+
+      #pragma clang loop unroll(full)
+      for(int j = 0; j < TN; j++) {
+        out_vec[out_col + j] = result[j];
+      }
+    }
+
+  }
+
+
+};
+
+///////////////////////////////////////////////////////////////////////////////
+/// Matrix vector multiplication
+///////////////////////////////////////////////////////////////////////////////
+
+template <
+    typename T, 
+    const int BM, /* Threadgroup rows (in threads) */
+    const int BN, /* Threadgroup cols (in threads) */
+    const int TM, /* Thread rows (in elements) */
+    const int TN> /* Thread cols (in elements) */
+[[kernel, max_total_threads_per_threadgroup(BM * BN)]] void gemv(
+    const device T* mat [[buffer(0)]],
+    const device T* in_vec [[buffer(1)]],
+    device T* out_vec [[buffer(2)]], 
+    const constant int& in_vec_size [[buffer(3)]],
+    const constant int& out_vec_size [[buffer(4)]],
+    const constant int& vector_batch_stride [[buffer(5)]],
+    const constant int& matrix_batch_stride [[buffer(6)]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint3 lid [[thread_position_in_threadgroup]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+
+  using gemv_kernel = GEMVKernel<T, BM, BN, TM, TN>;
+  threadgroup T tgp_memory[gemv_kernel::tgp_mem_size];
+
+  // Update batch offsets
+  in_vec += tid.z * vector_batch_stride;
+  mat += tid.z * matrix_batch_stride;
+  out_vec += tid.z * out_vec_size;
+
+  gemv_kernel::run( 
+    mat, 
+    in_vec, 
+    out_vec,
+    in_vec_size,
+    out_vec_size,
+    tgp_memory,
+    tid,
+    lid,
+    simd_gid,
+    simd_lid
+  );
+
 }
 
-#define instantiate_gemv(name, itype, bm, bn, tm, tn) \
+template <
+    typename T, 
+    const int BM, /* Threadgroup rows (in threads) */
+    const int BN, /* Threadgroup cols (in threads) */
+    const int TM, /* Thread rows (in elements) */
+    const int TN> /* Thread cols (in elements) */
+[[kernel, max_total_threads_per_threadgroup(BM * BN)]] void gemv_nc(
+    const device T* mat [[buffer(0)]],
+    const device T* in_vec [[buffer(1)]],
+    device T* out_vec [[buffer(2)]], 
+    const constant int& in_vec_size [[buffer(3)]],
+    const constant int& out_vec_size [[buffer(4)]],
+    const constant int& nc_dim [[buffer(5)]],
+    const device int* nc_shape [[buffer(6)]],
+    const device size_t* nc_strides_vec [[buffer(7)]],
+    const device size_t* nc_strides_mat [[buffer(8)]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint3 lid [[thread_position_in_threadgroup]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+
+  using gemv_kernel = GEMVKernel<T, BM, BN, TM, TN>;
+  threadgroup T tgp_memory[gemv_kernel::tgp_mem_size];
+
+  // Update batch offsets
+  in_vec += elem_to_loc(tid.z, nc_shape, nc_strides_vec, nc_dim);
+  mat += elem_to_loc(tid.z, nc_shape, nc_strides_mat, nc_dim);
+  out_vec += tid.z * out_vec_size;
+
+  gemv_kernel::run( 
+    mat, 
+    in_vec, 
+    out_vec,
+    in_vec_size,
+    out_vec_size,
+    tgp_memory,
+    tid,
+    lid,
+    simd_gid,
+    simd_lid
+  );
+
+}
+
+
+#define instantiate_gemv_c(name, itype, bm, bn, tm, tn) \
   template [[host_name("gemv_" #name "_bm" #bm "_bn" #bn "_tm" #tm "_tn" #tn)]] \
   [[kernel]] void gemv<itype, bm, bn, tm, tn>( \
     const device itype* mat [[buffer(0)]], \
@@ -145,28 +393,51 @@ template <typename T,
     const constant int& vector_batch_stride [[buffer(5)]], \
     const constant int& matrix_batch_stride [[buffer(6)]], \
     uint3 tid [[threadgroup_position_in_grid]], \
+    uint3 lid [[thread_position_in_threadgroup]], \
     uint simd_gid [[simdgroup_index_in_threadgroup]], \
     uint simd_lid [[thread_index_in_simdgroup]]);
 
-#define instantiate_gemv_blocks(name, itype) \
-    instantiate_gemv(name, itype, 4, 32, 1, 4) \
-    instantiate_gemv(name, itype, 4, 32, 4, 4) \
-    instantiate_gemv(name, itype, 8, 32, 4, 4)
+#define instantiate_gemv_nc(name, itype, bm, bn, tm, tn) \
+  template [[host_name("gemv_" #name "_bm" #bm "_bn" #bn "_tm" #tm "_tn" #tn "_nc")]] \
+  [[kernel]] void gemv_nc<itype, bm, bn, tm, tn>( \
+    const device itype* mat [[buffer(0)]], \
+    const device itype* vec [[buffer(1)]], \
+    device itype* out [[buffer(2)]], \
+    const constant int& in_vec_size [[buffer(3)]], \
+    const constant int& out_vec_size [[buffer(4)]], \
+    const constant int& nc_dim [[buffer(5)]], \
+    const device int* nc_shape [[buffer(6)]], \
+    const device size_t* nc_strides_vec [[buffer(7)]], \
+    const device size_t* nc_strides_mat [[buffer(8)]], \
+    uint3 tid [[threadgroup_position_in_grid]], \
+    uint3 lid [[thread_position_in_threadgroup]], \
+    uint simd_gid [[simdgroup_index_in_threadgroup]], \
+    uint simd_lid [[thread_index_in_simdgroup]]);
 
-instantiate_gemv_blocks(float32, float)
-instantiate_gemv_blocks(float16, half)
-instantiate_gemv_blocks(bfloat16, bfloat16_t)
+#define instantiate_gemv(name, itype, bm, bn, tm, tn) \
+  instantiate_gemv_c(name, itype, bm, bn, tm, tn) \
+  instantiate_gemv_nc(name, itype, bm, bn, tm, tn)
+
+#define instantiate_gemv_blocks(name, itype) \
+  instantiate_gemv(name, itype, 4, 32, 1, 4) \
+  instantiate_gemv(name, itype, 4, 32, 4, 4) \
+  instantiate_gemv(name, itype, 8, 32, 4, 4)
+
+instantiate_gemv_blocks(float32, float);
+instantiate_gemv_blocks(float16, half);
+instantiate_gemv_blocks(bfloat16, bfloat16_t);
 
 ///////////////////////////////////////////////////////////////////////////////
 /// Vector matrix multiplication
 ///////////////////////////////////////////////////////////////////////////////
 
-template <typename T, 
-          const int BM, /* Threadgroup rows (in threads) */
-          const int BN, /* Threadgroup cols (in threads) */
-          const int TM, /* Thread rows (in elements) */
-          const int TN> /* Thread cols (in elements) */
-[[kernel]] void gemv_t(
+template <
+    typename T, 
+    const int BM, /* Threadgroup rows (in threads) */
+    const int BN, /* Threadgroup cols (in threads) */
+    const int TM, /* Thread rows (in elements) */
+    const int TN> /* Thread cols (in elements) */
+[[kernel, max_total_threads_per_threadgroup(BM * BN)]] void gemv_t(
     const device T* mat [[buffer(0)]],
     const device T* in_vec [[buffer(1)]],
     device T* out_vec [[buffer(2)]], 
@@ -175,110 +446,77 @@ template <typename T,
     const constant int& vector_batch_stride [[buffer(5)]],
     const constant int& matrix_batch_stride [[buffer(6)]],
     uint3 tid [[threadgroup_position_in_grid]],
-    uint3 lid [[thread_position_in_threadgroup]]) {
+    uint3 lid [[thread_position_in_threadgroup]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
 
-    // - The matrix of size (M = in_vec_size, N = out_vec_size) is divided up 
-    //   into blocks of (BM * TM, BN * TN) divided amoung threadgroups
-    // - Every thread works on a block of (TM, TN)
-    // - We assume each thead group is launched with (BN, BM, 1) threads
-    //
-    // 1. A thread loads TN elements each from mat along TM contiguous rows 
-    //      and the corresponding scalar from the vector
-    // 2. The thread then multiplies and adds to accumulate its local result for the block
-    // 3. At the end, each thread has accumulated results over all blocks across the rows
-    //      These are then summed up across the threadgroup
-    // 4. Each threadgroup writes its accumulated BN * TN outputs
-    //
-    // Edge case handling:
-    // - The threadgroup with the largest tid will have blocks that exceed the matrix
-    //   * The blocks that start outside the matrix are never read (thread results remain zero)
-    //   * The last thread that partialy overlaps with the matrix is shifted inwards 
-    //     such that the thread block fits exactly in the matrix
+  using gemv_kernel = GEMVTKernel<T, BM, BN, TM, TN>;
+  threadgroup T tgp_memory[gemv_kernel::tgp_mem_size];
 
-    // Update batch offsets
-    in_vec += tid.z * vector_batch_stride;
-    mat += tid.z * matrix_batch_stride;
-    out_vec += tid.z * out_vec_size;
-    
-    // Thread local accumulation results
-    T result[TN] = {0};
-    T inter[TN];
-    T v_coeff[TM];
+  // Update batch offsets
+  in_vec += tid.z * vector_batch_stride;
+  mat += tid.z * matrix_batch_stride;
+  out_vec += tid.z * out_vec_size;
 
-    // Threadgroup accumulation results
-    threadgroup T tgp_results[BN][BM][TM];
+  gemv_kernel::run( 
+    mat, 
+    in_vec, 
+    out_vec,
+    in_vec_size,
+    out_vec_size,
+    tgp_memory,
+    tid,
+    lid,
+    simd_gid,
+    simd_lid
+  );
+}
 
-    int out_col = (tid.x * BN + lid.x) * TN;
-    int in_row = lid.y * TM;
+template <
+    typename T, 
+    const int BM, /* Threadgroup rows (in threads) */
+    const int BN, /* Threadgroup cols (in threads) */
+    const int TM, /* Thread rows (in elements) */
+    const int TN> /* Thread cols (in elements) */
+[[kernel, max_total_threads_per_threadgroup(BM * BN)]] void gemv_t_nc(
+    const device T* mat [[buffer(0)]],
+    const device T* in_vec [[buffer(1)]],
+    device T* out_vec [[buffer(2)]], 
+    const constant int& in_vec_size [[buffer(3)]],
+    const constant int& out_vec_size [[buffer(4)]],
+    const constant int& nc_dim [[buffer(5)]],
+    const device int* nc_shape [[buffer(6)]],
+    const device size_t* nc_strides_vec [[buffer(7)]],
+    const device size_t* nc_strides_mat [[buffer(8)]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint3 lid [[thread_position_in_threadgroup]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
 
-    // Edgecase handling
-    if (out_col < out_vec_size) {
-        // Edgecase handling
-        out_col = out_col + TN < out_vec_size ? out_col : out_vec_size - TN;
+  using gemv_kernel = GEMVTKernel<T, BM, BN, TM, TN>;
+  threadgroup T tgp_memory[gemv_kernel::tgp_mem_size];
 
-        // Per thread accumulation main loop
-        int bm = in_row;
-        for(; bm < in_vec_size; bm += BM * TM) {
-            // Adding a threadgroup_barrier improves performance slightly
-            // This is possibly it may help exploit cache better
-            threadgroup_barrier(mem_flags::mem_none);
+  // Update batch offsets
+  in_vec += elem_to_loc(tid.z, nc_shape, nc_strides_vec, nc_dim);
+  mat += elem_to_loc(tid.z, nc_shape, nc_strides_mat, nc_dim);
+  out_vec += tid.z * out_vec_size;
 
-            if(bm + TM <= in_vec_size) {
-
-                #pragma clang loop unroll(full)
-                for(int tm = 0; tm < TM; tm++) {
-                    v_coeff[tm] = in_vec[bm + tm];
-                }
-
-                #pragma clang loop unroll(full)
-                for(int tm = 0; tm < TM; tm++) {
-                    for(int tn = 0; tn < TN; tn++) {
-                        inter[tn] = mat[(bm + tm) * out_vec_size + out_col + tn];
-                    }
-                    for(int tn = 0; tn < TN; tn++) {
-                        result[tn] += v_coeff[tm] * inter[tn];
-                    }
-                }
-            
-            } else { // Edgecase handling
-                for(int tm = 0; bm + tm < in_vec_size; tm++) {
-                    v_coeff[tm] = in_vec[bm + tm];
-
-                    for(int tn = 0; tn < TN; tn++) {
-                        inter[tn] = mat[(bm + tm) * out_vec_size + out_col + tn];
-                    }
-                    for(int tn = 0; tn < TN; tn++) {
-                        result[tn] += v_coeff[tm] * inter[tn];
-                    }
-                }
-            }
-        }
-
-    }
-
-    // Threadgroup collection
-    for(int i = 0; i < TN; i++) {
-        tgp_results[lid.x][lid.y][i] = result[i];
-    }
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-
-    if(lid.y == 0 && out_col < out_vec_size) {
-        // Threadgroup accumulation
-        #pragma clang loop unroll(full)
-        for(int i = 1; i < BM; i++) {
-            for(int j = 0; j < TN; j++) {
-                result[j] += tgp_results[lid.x][i][j];
-            }
-        }
-
-        for(int j = 0; j < TN; j++) {
-            out_vec[out_col + j] = result[j];
-        }
-    }
+  gemv_kernel::run( 
+    mat, 
+    in_vec, 
+    out_vec,
+    in_vec_size,
+    out_vec_size,
+    tgp_memory,
+    tid,
+    lid,
+    simd_gid,
+    simd_lid
+  );
 
 }
 
-#define instantiate_gemv_t(name, itype, bm, bn, tm, tn) \
+#define instantiate_gemv_t_c(name, itype, bm, bn, tm, tn) \
   template [[host_name("gemv_t_" #name "_bm" #bm "_bn" #bn "_tm" #tm "_tn" #tn)]] \
   [[kernel]] void gemv_t<itype, bm, bn, tm, tn>( \
     const device itype* mat [[buffer(0)]], \
@@ -289,16 +527,39 @@ template <typename T,
     const constant int& vector_batch_stride [[buffer(5)]], \
     const constant int& matrix_batch_stride [[buffer(6)]], \
     uint3 tid [[threadgroup_position_in_grid]], \
-    uint3 lid [[thread_position_in_threadgroup]]);
+    uint3 lid [[thread_position_in_threadgroup]], \
+    uint simd_gid [[simdgroup_index_in_threadgroup]], \
+    uint simd_lid [[thread_index_in_simdgroup]]);
+
+#define instantiate_gemv_t_nc(name, itype, bm, bn, tm, tn) \
+  template [[host_name("gemv_t_" #name "_bm" #bm "_bn" #bn "_tm" #tm "_tn" #tn "_nc")]] \
+  [[kernel]] void gemv_t_nc<itype, bm, bn, tm, tn>( \
+    const device itype* mat [[buffer(0)]], \
+    const device itype* vec [[buffer(1)]], \
+    device itype* out [[buffer(2)]], \
+    const constant int& in_vec_size [[buffer(3)]], \
+    const constant int& out_vec_size [[buffer(4)]], \
+    const constant int& nc_dim [[buffer(5)]], \
+    const device int* nc_shape [[buffer(6)]], \
+    const device size_t* nc_strides_vec [[buffer(7)]], \
+    const device size_t* nc_strides_mat [[buffer(8)]], \
+    uint3 tid [[threadgroup_position_in_grid]], \
+    uint3 lid [[thread_position_in_threadgroup]], \
+    uint simd_gid [[simdgroup_index_in_threadgroup]], \
+    uint simd_lid [[thread_index_in_simdgroup]]);
+
+#define instantiate_gemv_t(name, itype, bm, bn, tm, tn) \
+  instantiate_gemv_t_c(name, itype, bm, bn, tm, tn) \
+  instantiate_gemv_t_nc(name, itype, bm, bn, tm, tn)
 
 #define instantiate_gemv_t_blocks(name, itype) \
-    instantiate_gemv_t(name, itype, 8, 8, 4, 1) \
-    instantiate_gemv_t(name, itype, 8, 8, 4, 4) \
-    instantiate_gemv_t(name, itype, 8, 16, 4, 4) \
-    instantiate_gemv_t(name, itype, 8, 32, 4, 4) \
-    instantiate_gemv_t(name, itype, 8, 64, 4, 4) \
-    instantiate_gemv_t(name, itype, 8, 128, 4, 4)
+  instantiate_gemv_t(name, itype, 8, 8, 4, 1) \
+  instantiate_gemv_t(name, itype, 8, 8, 4, 4) \
+  instantiate_gemv_t(name, itype, 8, 16, 4, 4) \
+  instantiate_gemv_t(name, itype, 8, 32, 4, 4) \
+  instantiate_gemv_t(name, itype, 8, 64, 4, 4) \
+  instantiate_gemv_t(name, itype, 8, 128, 4, 4)
 
-instantiate_gemv_t_blocks(float32, float)
-instantiate_gemv_t_blocks(float16, half)
-instantiate_gemv_t_blocks(bfloat16, bfloat16_t)
+instantiate_gemv_t_blocks(float32, float);
+instantiate_gemv_t_blocks(float16, half);
+instantiate_gemv_t_blocks(bfloat16, bfloat16_t);

mlx/backend/metal/kernels/quantized.metal

@@ -0,0 +1,287 @@
+// Copyright © 2023 Apple Inc.
+
+#include <metal_stdlib>
+#include <metal_simdgroup>
+
+#include "mlx/backend/metal/kernels/bf16.h"
+#include "mlx/backend/metal/kernels/defines.h"
+#include "mlx/backend/metal/kernels/gemm/gemm.h"
+#include "mlx/backend/metal/kernels/utils.h"
+
+using namespace metal;
+
+#define MLX_MTL_CONST static constant constexpr const
+
+MLX_MTL_CONST int SIMD_SIZE = 32;
+
+template <typename T, const int BM, const int BN, const int group_size, const int bits>
+[[kernel]] void qmv(
+    const device uint32_t* w [[buffer(0)]],
+    const device T* scales [[buffer(1)]],
+    const device T* biases [[buffer(2)]],
+    const device T* x [[buffer(3)]],
+    device T* y [[buffer(4)]],
+    const constant int& in_vec_size [[buffer(5)]],
+    const constant int& out_vec_size [[buffer(6)]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint lid [[thread_index_in_threadgroup]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+
+  static_assert(BN == SIMD_SIZE, "qmv expects BN to be equal to SIMD_SIZE");
+
+  constexpr int bitmask = (1 << bits) - 1;
+  constexpr int el_per_thread = 32 / bits;
+  constexpr int colgroup = BN * el_per_thread;
+  constexpr int groups_per_block = colgroup / group_size;
+  constexpr int simdgroups_fetching_vec = colgroup / SIMD_SIZE;
+
+  threadgroup T scales_block[BM * groups_per_block];
+  threadgroup T biases_block[BM * groups_per_block];
+  threadgroup T x_block[colgroup];
+
+  thread uint32_t w_local;
+  thread T result = 0;
+  thread T scale = 1;
+  thread T bias = 0;
+  thread T x_thread[el_per_thread];
+
+  // Adjust positions
+  const int in_vec_size_w = in_vec_size / el_per_thread;
+  const int in_vec_size_g = in_vec_size / group_size;
+  int out_row = tid.y * BM + simd_gid;
+  w += out_row * in_vec_size_w;
+  scales += out_row * in_vec_size_g;
+  biases += out_row * in_vec_size_g;
+  x += tid.z * in_vec_size;
+  y += tid.z * out_vec_size;
+
+  // Loop over in_vec in blocks of colgroup
+  for (int i=0; i<in_vec_size; i+=colgroup) {
+    // Load the vec to shared memory
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    if (simd_gid < simdgroups_fetching_vec) {
+      x_block[lid] = x[lid + i];
+    }
+    if (simd_lid == 0) {
+      #pragma clang loop unroll(full)
+      for (int j=0; j<groups_per_block; j++) {
+        scales_block[simd_gid * groups_per_block + j] = scales[i / group_size + j];
+      }
+      #pragma clang loop unroll(full)
+      for (int j=0; j<groups_per_block; j++) {
+        biases_block[simd_gid * groups_per_block + j] = biases[i / group_size + j];
+      }
+    }
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    // Load in_vec, scale, bias to registers
+    #pragma clang loop unroll(full)
+    for (int j=0; j<el_per_thread; j++) {
+      x_thread[j] = x_block[simd_lid*el_per_thread + j];
+    }
+    scale = scales_block[simd_gid * groups_per_block + simd_lid * el_per_thread / group_size];
+    bias = biases_block[simd_gid * groups_per_block + simd_lid * el_per_thread / group_size];
+
+    // Load the matrix elements
+    w_local = w[i / el_per_thread + simd_lid];
+
+    // Do all the work.
+    #pragma clang loop unroll(full)
+    for (int k=0; k<el_per_thread; k++) {
+      result += (scale * static_cast<T>(w_local & bitmask) + bias) * x_thread[k];
+      w_local >>= bits;
+    }
+  }
+
+  // Accumulate in the simdgroup
+  result = simd_sum(result);
+
+  // Store the result
+  if (simd_lid == 0) {
+    y[out_row] = result;
+  }
+}
+
+
+template <typename T, const int BM, const int BK, const int BN, const int group_size, const int bits>
+[[kernel]] void qmm_t(
+    const device T* x [[buffer(0)]],
+    const device uint32_t* w [[buffer(1)]],
+    const device T* scales [[buffer(2)]],
+    const device T* biases [[buffer(3)]],
+    device T* y [[buffer(4)]],
+    const constant int& M [[buffer(5)]],
+    const constant int& N [[buffer(6)]],
+    const constant int& K [[buffer(7)]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint lid [[thread_index_in_threadgroup]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+
+  static_assert(BK >= SIMD_SIZE, "BK should be larger than SIMD_SIZE");
+  static_assert(BK % SIMD_SIZE == 0, "BK should be divisible by SIMD_SIZE");
+
+  const uint lidy = lid / SIMD_SIZE;
+
+  constexpr int WM = 2;
+  constexpr int WN = 2;
+  constexpr int bitmask = (1 << bits) - 1;
+  constexpr int el_per_int = 32 / bits;
+  constexpr int ints_per_block = BK / el_per_int;
+  constexpr int groups_per_block = (BK / group_size > 0) ? (BK / group_size) : 1;
+  constexpr int groups_per_simd = BN / (WM * WN);
+  constexpr int w_els_per_thread = (BN * BK / el_per_int) / (SIMD_SIZE * WM * WN);
+
+  // Using the kernel just as a type to instantiate the appropriate BlockMMA
+  // and constexpr size calculations
+  using mma_t = BlockMMA<T, BM, BN, BK, WM, WN, false, true>;
+  using loader_x_t = BlockLoader<T, BM, BK, BK, 4, WM * WN * SIMD_SIZE, false, true, 0>;
+
+  threadgroup T scales_block[BN * groups_per_block];
+  threadgroup T biases_block[BN * groups_per_block];
+  threadgroup T Xs[BM * BK];
+  threadgroup T Ws[BN * BK];
+
+  // Set the block
+  const int K_w = K / el_per_int;
+  const int K_g = K / group_size;
+  const int y_row = tid.y * BM;
+  const int y_col = tid.x * BN;
+  x += y_row * K;
+  w += y_col * K_w;
+  scales += y_col * K_g;
+  biases += y_col * K_g;
+  y += y_row * N + y_col;
+
+  // Make the x loader and mma operation
+  const short num_els = min(BM, M - y_row);
+  loader_x_t loader_x(x, K, Xs, simd_gid, simd_lid);
+  mma_t mma_op(simd_gid, simd_lid);
+
+  for (int k=0; k<K; k += BK) {
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    // Load the x tile
+    if (num_els < BM) {
+        loader_x.load_safe(short2(BK, num_els));
+    } else {
+        loader_x.load_unsafe();
+    }
+
+    // Load the scale and bias
+    if (simd_lid == 0) {
+      threadgroup T *scales_block_local = scales_block + lidy * groups_per_block * groups_per_simd;
+      threadgroup T *biases_block_local = biases_block + lidy * groups_per_block * groups_per_simd;
+      const device T *scales_local = scales + lidy * groups_per_simd * K_g + k / group_size;
+      const device T *biases_local = biases + lidy * groups_per_simd * K_g + k / group_size;
+      #pragma clang loop unroll(full)
+      for (int gs=0; gs<groups_per_simd; gs++) {
+        #pragma clang loop unroll(full)
+        for (int gc=0; gc<groups_per_block; gc++) {
+          scales_block_local[gc] = scales_local[gc];
+          biases_block_local[gc] = biases_local[gc];
+        }
+        scales_block_local += groups_per_block;
+        scales_local += K_g;
+        biases_block_local += groups_per_block;
+        biases_local += K_g;
+      }
+    }
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    // Load the w tile
+    {
+      for (int wo=0; wo<w_els_per_thread; wo++) {
+        int offset = lid * w_els_per_thread + wo;
+        int offset_row = offset / (BK / el_per_int);
+        int offset_col = offset % (BK / el_per_int);
+        const device uint32_t * w_local = w + offset_row * K_w + offset_col;
+        threadgroup T * Ws_local = Ws + offset_row * BK + offset_col * el_per_int;
+
+        uint32_t wi = *w_local;
+        T scale = scales_block[offset_row * groups_per_block + offset_col / (group_size / el_per_int)];
+        T bias = biases_block[offset_row * groups_per_block + offset_col / (group_size / el_per_int)];
+
+        #pragma clang loop unroll(full)
+        for (int t=0; t<el_per_int; t++) {
+          Ws_local[t] = scale * static_cast<T>(wi & bitmask) + bias;
+          wi >>= bits;
+        }
+      }
+    }
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    // Multiply and accumulate threadgroup elements
+    mma_op.mma(Xs, Ws);
+
+    // Prepare for next iteration
+    loader_x.next();
+    w += ints_per_block;
+    // scales and biases cannot be advanced because they would have to be
+    // advanced every other iteration or sth.
+  }
+
+  // Store results to device memory
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  if (num_els < BM) {
+    mma_op.store_result_safe(y, N, short2(BN, num_els));
+  } else {
+    mma_op.store_result(y, N);
+  }
+}
+
+
+#define instantiate_qmv(name, itype, group_size, bits) \
+  template [[host_name("qmv_n_" #name "_gs_" #group_size "_b_" #bits)]] \
+  [[kernel]] void qmv<itype, 32, 32, group_size, bits>( \
+    const device uint32_t* w [[buffer(0)]], \
+    const device itype* scales [[buffer(1)]], \
+    const device itype* biases [[buffer(2)]], \
+    const device itype* x [[buffer(3)]], \
+    device itype* y [[buffer(4)]], \
+    const constant int& in_vec_size [[buffer(5)]], \
+    const constant int& out_vec_size [[buffer(6)]], \
+    uint3 tid [[threadgroup_position_in_grid]], \
+    uint lid [[thread_index_in_threadgroup]], \
+    uint simd_gid [[simdgroup_index_in_threadgroup]], \
+    uint simd_lid [[thread_index_in_simdgroup]]);
+
+#define instantiate_qmv_types(group_size, bits) \
+  instantiate_qmv(float32, float, group_size, bits) \
+  instantiate_qmv(float16, half, group_size, bits) \
+  instantiate_qmv(bfloat16, bfloat16_t, group_size, bits)
+
+instantiate_qmv_types(128, 2)
+instantiate_qmv_types(128, 4)
+instantiate_qmv_types(128, 8)
+instantiate_qmv_types( 64, 2)
+instantiate_qmv_types( 64, 4)
+instantiate_qmv_types( 64, 8)
+
+#define instantiate_qmm_t(name, itype, group_size, bits) \
+  template [[host_name("qmm_t_" #name "_gs_" #group_size "_b_" #bits)]] \
+  [[kernel]] void qmm_t<itype, 32, 64, 32, group_size, bits>( \
+      const device itype* x [[buffer(0)]], \
+      const device uint32_t* w [[buffer(1)]], \
+      const device itype* scales [[buffer(2)]], \
+      const device itype* biases [[buffer(3)]], \
+      device itype* y [[buffer(4)]], \
+      const constant int& M [[buffer(5)]], \
+      const constant int& N [[buffer(6)]], \
+      const constant int& K [[buffer(7)]], \
+      uint3 tid [[threadgroup_position_in_grid]], \
+      uint lid [[thread_index_in_threadgroup]], \
+      uint simd_gid [[simdgroup_index_in_threadgroup]], \
+      uint simd_lid [[thread_index_in_simdgroup]]);
+
+#define instantiate_qmm_t_types(group_size, bits) \
+  instantiate_qmm_t(float32, float, group_size, bits) \
+  instantiate_qmm_t(float16, half, group_size, bits) \
+  instantiate_qmm_t(bfloat16, bfloat16_t, group_size, bits)
+
+instantiate_qmm_t_types(128, 2)
+instantiate_qmm_t_types(128, 4)
+instantiate_qmm_t_types(128, 8)
+instantiate_qmm_t_types( 64, 2)
+instantiate_qmm_t_types( 64, 4)
+instantiate_qmm_t_types( 64, 8)

mlx/backend/metal/kernels/sort.metal

@@ -9,7 +9,7 @@
 #define MLX_MTL_CONST static constant constexpr const
 #define MLX_MTL_LOOP_UNROLL _Pragma("clang loop unroll(full)")
 
-using namespace metal;\
+using namespace metal;
 
 // Based on GPU merge sort algorithm at https://github.com/NVIDIA/cccl/tree/main/cub/cub
 

mlx/backend/metal/kernels/unary.metal

@@ -43,6 +43,19 @@ struct ArcTanh {
   template <typename T> T operator()(T x) { return metal::precise::atanh(x); };
 };
 
+struct Ceil {
+  template <typename T> T operator()(T x) { return metal::ceil(x); };
+  template <> int8_t operator()(int8_t x) { return x; };
+  template <> int16_t operator()(int16_t x) { return x; };
+  template <> int32_t operator()(int32_t x) { return x; };
+  template <> int64_t operator()(int64_t x) { return x; };
+  template <> uint8_t operator()(uint8_t x) { return x; };
+  template <> uint16_t operator()(uint16_t x) { return x; };
+  template <> uint32_t operator()(uint32_t x) { return x; };
+  template <> uint64_t operator()(uint64_t x) { return x; };
+  template <> bool operator()(bool x) { return x; };
+};
+
 struct Cos {
   template <typename T> T operator()(T x) { return metal::precise::cos(x); };
 
@@ -83,6 +96,19 @@ struct Exp {
   }
 };
 
+struct Floor {
+  template <typename T> T operator()(T x) { return metal::floor(x); };
+  template <> int8_t operator()(int8_t x) { return x; };
+  template <> int16_t operator()(int16_t x) { return x; };
+  template <> int32_t operator()(int32_t x) { return x; };
+  template <> int64_t operator()(int64_t x) { return x; };
+  template <> uint8_t operator()(uint8_t x) { return x; };
+  template <> uint16_t operator()(uint16_t x) { return x; };
+  template <> uint32_t operator()(uint32_t x) { return x; };
+  template <> uint64_t operator()(uint64_t x) { return x; };
+  template <> bool operator()(bool x) { return x; };
+};
+
 struct Log {
   template <typename T> T operator()(T x) { return metal::precise::log(x); };
 };
@@ -107,6 +133,11 @@ struct Negative {
   template <typename T> T operator()(T x) { return -x; };
 };
 
+struct Round {
+  template <typename T> T operator()(T x) { return metal::round(x); };
+  template <> complex64_t operator()(complex64_t x) { return {metal::round(x.real), metal::round(x.imag)}; };
+};
+
 struct Sigmoid {
   template <typename T>
   T operator()(T x) {
@@ -253,9 +284,11 @@ instantiate_unary_float(arcsin, ArcSin)
 instantiate_unary_float(arcsinh, ArcSinh)
 instantiate_unary_float(arctan, ArcTan)
 instantiate_unary_float(arctanh, ArcTanh)
+instantiate_unary_types(ceil, Ceil)
 instantiate_unary_float(cos, Cos)
 instantiate_unary_float(cosh, Cosh)
 instantiate_unary_float(exp, Exp)
+instantiate_unary_types(floor, Floor)
 instantiate_unary_float(log, Log)
 instantiate_unary_float(log2, Log2)
 instantiate_unary_float(log10, Log10)
@@ -272,6 +305,7 @@ instantiate_unary_float(sqrt, Sqrt)
 instantiate_unary_float(rsqrt, Rsqrt)
 instantiate_unary_float(tan, Tan)
 instantiate_unary_float(tanh, Tanh)
+instantiate_unary_float(round, Round)
 
 instantiate_unary_all(abs, complex64, complex64_t, Abs)
 instantiate_unary_all(cos, complex64, complex64_t, Cos)
@@ -282,5 +316,6 @@ instantiate_unary_all(sin, complex64, complex64_t, Sin)
 instantiate_unary_all(sinh, complex64, complex64_t, Sinh)
 instantiate_unary_all(tan, complex64, complex64_t, Tan)
 instantiate_unary_all(tanh, complex64, complex64_t, Tanh)
+instantiate_unary_all(round, complex64, complex64_t, Round)
 
 instantiate_unary_all(lnot, bool_, bool, LogicalNot)

mlx/backend/metal/matmul.cpp

@@ -343,10 +343,18 @@ void Matmul::eval_gpu(const std::vector<array>& inputs, array& out) {
     int mat_rows = transpose_mat ? in_vector_len : out_vector_len;
 
     int batch_size_mat = mat.data_size() / (mat_cols * mat_rows);
-    int stride_mat = batch_size_mat == batch_size_out ? mat_cols * mat_rows : 0;
+    int stride_mat = batch_size_mat == 1 ? 0 : mat_cols * mat_rows;
 
     int batch_size_vec = vec.data_size() / in_vector_len;
-    int stride_vec = batch_size_vec == batch_size_out ? in_vector_len : 0;
+    int stride_vec = batch_size_vec == 1 ? 0 : in_vector_len;
+
+    // Determine if inputs have simple batching / broadcasting
+    bool contiguous_kernel =
+        (batch_size_out == std::max(batch_size_mat, batch_size_vec) &&
+         (batch_size_mat == batch_size_vec ||
+          std::min(batch_size_mat, batch_size_vec) == 1));
+
+    int nc_dim = out.ndim() - 2;
 
     // Determine dispatch kernel
     int tm = 4, tn = 4;
@@ -383,6 +391,10 @@ void Matmul::eval_gpu(const std::vector<array>& inputs, array& out) {
 
     kname << "_bm" << bm << "_bn" << bn << "_tm" << tm << "_tn" << tn;
 
+    if (!contiguous_kernel) {
+      kname << "_nc";
+    }
+
     // Encode and dispatch kernel
     auto compute_encoder = d.get_command_encoder(s.index);
     auto kernel = d.get_kernel(kname.str());
@@ -398,8 +410,22 @@ void Matmul::eval_gpu(const std::vector<array>& inputs, array& out) {
 
     compute_encoder->setBytes(&in_vector_len, sizeof(int), 3);
     compute_encoder->setBytes(&out_vector_len, sizeof(int), 4);
-    compute_encoder->setBytes(&stride_vec, sizeof(int), 5);
-    compute_encoder->setBytes(&stride_mat, sizeof(int), 6);
+
+    if (contiguous_kernel) {
+      compute_encoder->setBytes(&stride_vec, sizeof(int), 5);
+      compute_encoder->setBytes(&stride_mat, sizeof(int), 6);
+    } else {
+      // In case of complex broadcasting, we consider the shape[:-2] and
+      // strides [:-2] to determine the location of a batch
+      // nc_dim = out.ndim() - 2
+      compute_encoder->setBytes(&nc_dim, sizeof(int), 5);
+      compute_encoder->setBytes(out.shape().data(), nc_dim * sizeof(int), 6);
+      compute_encoder->setBytes(
+          vec.strides().data(), nc_dim * sizeof(size_t), 7);
+      compute_encoder->setBytes(
+          mat.strides().data(), nc_dim * sizeof(size_t), 8);
+    }
+
     compute_encoder->dispatchThreadgroups(grid_dims, group_dims);
 
     d.get_command_buffer(s.index)->addCompletedHandler(

mlx/backend/metal/primitives.cpp

@@ -84,9 +84,9 @@ void binary_op(
     }
 
     // Launch up to 3D grid of threads
-    int dim0 = ndim > 0 ? shape[ndim - 1] : 1;
-    int dim1 = ndim > 1 ? shape[ndim - 2] : 1;
-    int rest = out.size() / (dim0 * dim1);
+    size_t dim0 = ndim > 0 ? shape[ndim - 1] : 1;
+    size_t dim1 = ndim > 1 ? shape[ndim - 2] : 1;
+    size_t rest = out.size() / (dim0 * dim1);
     NS::UInteger thread_group_size = kernel->maxTotalThreadsPerThreadgroup();
     if (thread_group_size != 1024) {
       throw std::runtime_error("[Metal::binary] Must use 1024 sized block");
@@ -215,7 +215,8 @@ void Arange::eval_gpu(const std::vector<array>& inputs, array& out) {
       arange_set_scalars<float>(start_, start_ + step_, compute_encoder);
       break;
     case bfloat16:
-      throw std::runtime_error("[Arange::eval_gpu] Does not support bfloat16");
+      arange_set_scalars<bfloat16_t>(start_, start_ + step_, compute_encoder);
+      break;
     case complex64:
       throw std::runtime_error("[Arange::eval_gpu] Does not support complex64");
   }
@@ -363,6 +364,10 @@ void Divide::eval_gpu(const std::vector<array>& inputs, array& out) {
   binary_op(inputs, out, "div");
 }
 
+void Remainder::eval_gpu(const std::vector<array>& inputs, array& out) {
+  binary_op(inputs, out, "rem");
+}
+
 void Equal::eval_gpu(const std::vector<array>& inputs, array& out) {
   binary_op(inputs, out, equal_nan_ ? "naneq" : "eq");
 }
@@ -446,6 +451,14 @@ void Minimum::eval_gpu(const std::vector<array>& inputs, array& out) {
   binary_op(inputs, out, "min");
 }
 
+void Floor::eval_gpu(const std::vector<array>& inputs, array& out) {
+  unary_op(inputs, out, "floor");
+}
+
+void Ceil::eval_gpu(const std::vector<array>& inputs, array& out) {
+  unary_op(inputs, out, "ceil");
+}
+
 void Multiply::eval_gpu(const std::vector<array>& inputs, array& out) {
   binary_op(inputs, out, "mul");
 }
@@ -551,6 +564,17 @@ void Reshape::eval_gpu(const std::vector<array>& inputs, array& out) {
   }
 }
 
+void Round::eval_gpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  const auto& in = inputs[0];
+  if (not is_integral(in.dtype())) {
+    unary_op(inputs, out, "round");
+  } else {
+    // No-op integer types
+    out.copy_shared_buffer(in);
+  }
+}
+
 void Sigmoid::eval_gpu(const std::vector<array>& inputs, array& out) {
   unary_op(inputs, out, "sigmoid");
 }

mlx/backend/metal/quantized.cpp

@@ -0,0 +1,123 @@
+// Copyright © 2023 Apple Inc.
+
+#include <cassert>
+#include <iostream>
+
+#include "mlx/backend/metal/copy.h"
+#include "mlx/backend/metal/device.h"
+#include "mlx/backend/metal/utils.h"
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+
+void QuantizedMatmul::eval_gpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 4);
+
+  out.set_data(allocator::malloc_or_wait(out.nbytes()));
+  auto& s = stream();
+  auto& d = metal::device(s.device);
+
+  auto& x_pre = inputs[0];
+  auto& w_pre = inputs[1];
+  auto& scales_pre = inputs[2];
+  auto& biases_pre = inputs[3];
+
+  std::vector<array> copies;
+  auto check_transpose = [&copies, &s](const array& arr) {
+    auto stx = arr.strides()[arr.ndim() - 2];
+    auto sty = arr.strides()[arr.ndim() - 1];
+    if (stx == arr.shape(-1) && sty == 1) {
+      return std::make_tuple(false, stx, arr);
+    } else if (stx == 1 && sty == arr.shape(-2)) {
+      return std::make_tuple(true, sty, arr);
+    } else {
+      array arr_copy(arr.shape(), arr.dtype(), nullptr, {});
+      copy_gpu(arr, arr_copy, CopyType::General, s);
+      copies.push_back(arr_copy);
+      size_t stx = arr.shape(-1);
+      return std::make_tuple(false, stx, arr_copy);
+    }
+  };
+  auto [x_transposed, x_cols, x] = check_transpose(x_pre);
+  auto [w_transposed, w_cols, w] = check_transpose(w_pre);
+  auto [scales_transposed, scales_cols, scales] = check_transpose(scales_pre);
+  auto [biases_transposed, biases_cols, biases] = check_transpose(biases_pre);
+
+  if (!w_transposed) {
+    throw std::runtime_error("The quantized weight should be transposed.");
+  }
+
+  if (x_transposed || scales_transposed || biases_transposed) {
+    throw std::runtime_error("x, scales and biases should be row contiguous.");
+  }
+
+  int D = x.shape(-1);
+  int B = x.size() / D;
+
+  // Route to the qmv kernel
+  if (B == 1) {
+    std::ostringstream kname;
+    kname << "qmv_" << (w_transposed ? "n_" : "t_") << type_to_name(out)
+          << "_gs_" << group_size_ << "_b_" << bits_;
+
+    // Encode and dispatch kernel
+    auto compute_encoder = d.get_command_encoder(s.index);
+    auto kernel = d.get_kernel(kname.str());
+    compute_encoder->setComputePipelineState(kernel);
+
+    int O = w.size() / w_cols;
+
+    int bo = 32;
+    int bd = 32;
+    MTL::Size group_dims = MTL::Size(bd, bo, 1);
+    MTL::Size grid_dims = MTL::Size(1, O / bo, B);
+
+    set_array_buffer(compute_encoder, w, 0);
+    set_array_buffer(compute_encoder, scales, 1);
+    set_array_buffer(compute_encoder, biases, 2);
+    set_array_buffer(compute_encoder, x, 3);
+    set_array_buffer(compute_encoder, out, 4);
+    compute_encoder->setBytes(&D, sizeof(int), 5);
+    compute_encoder->setBytes(&O, sizeof(int), 6);
+
+    compute_encoder->dispatchThreadgroups(grid_dims, group_dims);
+  }
+
+  // Route to the qmm kernel
+  else {
+    std::ostringstream kname;
+    kname << "qmm_" << (w_transposed ? "t_" : "n_") << type_to_name(out)
+          << "_gs_" << group_size_ << "_b_" << bits_;
+
+    // Encode and dispatch kernel
+    auto compute_encoder = d.get_command_encoder(s.index);
+    auto kernel = d.get_kernel(kname.str());
+    compute_encoder->setComputePipelineState(kernel);
+
+    int O = w.size() / w_cols;
+
+    int wn = 2;
+    int wm = 2;
+    int bm = 32;
+    int bn = 32;
+    int bk = 64;
+    MTL::Size group_dims = MTL::Size(32, wn, wm);
+    MTL::Size grid_dims = MTL::Size(O / bn, (B + bm - 1) / bm, 1);
+
+    set_array_buffer(compute_encoder, x, 0);
+    set_array_buffer(compute_encoder, w, 1);
+    set_array_buffer(compute_encoder, scales, 2);
+    set_array_buffer(compute_encoder, biases, 3);
+    set_array_buffer(compute_encoder, out, 4);
+    compute_encoder->setBytes(&B, sizeof(int), 5);
+    compute_encoder->setBytes(&O, sizeof(int), 6);
+    compute_encoder->setBytes(&D, sizeof(int), 7);
+
+    compute_encoder->dispatchThreadgroups(grid_dims, group_dims);
+  }
+
+  d.get_command_buffer(s.index)->addCompletedHandler(
+      [copies](MTL::CommandBuffer*) mutable { copies.clear(); });
+}
+
+} // namespace mlx::core

mlx/backend/no_metal/primitives.cpp

@@ -24,17 +24,20 @@ NO_GPU(ArgSort)
 NO_GPU(AsType)
 NO_GPU(AsStrided)
 NO_GPU(Broadcast)
+NO_GPU(Ceil)
 NO_GPU(Concatenate)
 NO_GPU(Convolution)
 NO_GPU(Copy)
 NO_GPU(Cos)
 NO_GPU(Cosh)
 NO_GPU(Divide)
+NO_GPU(Remainder)
 NO_GPU(Equal)
 NO_GPU(Erf)
 NO_GPU(ErfInv)
 NO_GPU(Exp)
 NO_GPU(FFT)
+NO_GPU(Floor)
 NO_GPU(Full)
 NO_GPU(Gather)
 NO_GPU(Greater)
@@ -55,9 +58,11 @@ NO_GPU(NotEqual)
 NO_GPU(Pad)
 NO_GPU(Partition)
 NO_GPU(Power)
+NO_GPU(QuantizedMatmul)
 NO_GPU(RandomBits)
 NO_GPU(Reduce)
 NO_GPU(Reshape)
+NO_GPU(Round)
 NO_GPU(Scan)
 NO_GPU(Scatter)
 NO_GPU(Sigmoid)

mlx/dtype.h

@@ -84,6 +84,10 @@ inline bool is_floating_point(const Dtype& t) {
       kindof(t) == Dtype::Kind::c;
 }
 
+inline bool is_complex(const Dtype& t) {
+  return kindof(t) == Dtype::Kind::c;
+}
+
 inline bool is_integral(const Dtype& t) {
   return !(is_floating_point(t));
 }

mlx/ops.cpp

@@ -129,6 +129,27 @@ array arange(int stop, StreamOrDevice s /* = {} */) {
   return arange(0.0, static_cast<double>(stop), 1.0, int32, to_stream(s));
 }
 
+array linspace(
+    double start,
+    double stop,
+    int num /* = 50 */,
+    Dtype dtype /* = float32 */,
+    StreamOrDevice s /* = {} */) {
+  if (num < 0) {
+    std::ostringstream msg;
+    msg << "[linspace] number of samples, " << num << ", must be non-negative.";
+    throw std::invalid_argument(msg.str());
+  }
+  array sequence = arange(0, num, float32, to_stream(s));
+  float step = (stop - start) / (num - 1);
+  return astype(
+      add(multiply(sequence, array(step), to_stream(s)),
+          array(start),
+          to_stream(s)),
+      dtype,
+      to_stream(s));
+}
+
 array astype(const array& a, Dtype dtype, StreamOrDevice s /* = {} */) {
   if (dtype == a.dtype()) {
     return a;
@@ -194,6 +215,52 @@ array ones_like(const array& a, StreamOrDevice s /* = {} */) {
   return ones(a.shape(), a.dtype(), to_stream(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.");
+  }
+  array result = zeros({n * m}, dtype, s);
+  if (k >= m || -k >= n) {
+    return reshape(result, {n, m}, s);
+  }
+
+  int diagonal_length = k >= 0 ? std::min(n, m - k) : std::min(n + k, m);
+  int start_index = (k >= 0) ? k : -k * m;
+
+  array diag_indices_array = arange(
+      start_index, start_index + diagonal_length * (m + 1), m + 1, int32, s);
+  array ones_array = ones({diagonal_length, 1}, dtype, s);
+  result = scatter(result, diag_indices_array, ones_array, 0, s);
+
+  return reshape(result, {n, m}, s);
+}
+
+array identity(int n, Dtype dtype, StreamOrDevice s /* = {} */) {
+  return eye(n, n, 0, dtype, s);
+}
+
+array tri(int n, int m, int k, Dtype type, StreamOrDevice s /* = {} */) {
+  auto l = expand_dims(arange(n, s), 1, s);
+  auto r = expand_dims(arange(-k, m - k, s), 0, s);
+  return astype(greater_equal(l, r, s), type, s);
+}
+
+array tril(array x, int k, StreamOrDevice s /* = {} */) {
+  if (x.ndim() < 2) {
+    throw std::invalid_argument("[tril] array must be atleast 2-D");
+  }
+  auto mask = tri(x.shape(-2), x.shape(-1), k, x.dtype(), s);
+  return where(mask, x, zeros_like(x, s), s);
+}
+
+array triu(array x, int k, StreamOrDevice s /* = {} */) {
+  if (x.ndim() < 2) {
+    throw std::invalid_argument("[triu] array must be atleast 2-D");
+  }
+  auto mask = tri(x.shape(-2), x.shape(-1), k - 1, x.dtype(), s);
+  return where(mask, zeros_like(x, s), x, s);
+}
+
 array reshape(
     const array& a,
     std::vector<int> shape,
@@ -231,6 +298,49 @@ array reshape(
       shape, a.dtype(), std::make_unique<Reshape>(to_stream(s), shape), {a});
 }
 
+array flatten(
+    const array& a,
+    int start_axis,
+    int end_axis /* = -1 */,
+    StreamOrDevice s /* = {} */) {
+  auto ndim = static_cast<int>(a.ndim());
+  auto start_ax = start_axis + (start_axis < 0 ? ndim : 0);
+  auto end_ax = end_axis + (end_axis < 0 ? ndim : 0);
+  start_ax = std::max(0, start_ax);
+  end_ax = std::min(ndim - 1, end_ax);
+  if (a.ndim() == 0) {
+    return reshape(a, {1}, s);
+  }
+  if (end_ax < start_ax) {
+    throw std::invalid_argument(
+        "[flatten] start_axis must be less than or equal to end_axis");
+  }
+  if (start_ax >= ndim) {
+    std::ostringstream msg;
+    msg << "[flatten] Invalid start_axis " << start_axis << " for array with "
+        << ndim << " dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+  if (end_ax < 0) {
+    std::ostringstream msg;
+    msg << "[flatten] Invalid end_axis " << end_axis << " for array with "
+        << ndim << " dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+  if (start_ax == end_ax) {
+    return a;
+  }
+  std::vector<int> new_shape(a.shape().begin(), a.shape().begin() + start_ax);
+  new_shape.push_back(-1);
+  new_shape.insert(
+      new_shape.end(), a.shape().begin() + end_ax + 1, a.shape().end());
+  return reshape(a, new_shape, s);
+}
+
+array flatten(const array& a, StreamOrDevice s /* = {} */) {
+  return flatten(a, 0, a.ndim() - 1, s);
+}
+
 array squeeze(
     const array& a,
     const std::vector<int>& axes,
@@ -506,6 +616,24 @@ split(const array& a, int num_splits, StreamOrDevice s /* = {} */) {
   return split(a, num_splits, 0, to_stream(s));
 }
 
+array clip(
+    const array& a,
+    const std::optional<array>& a_min,
+    const std::optional<array>& a_max,
+    StreamOrDevice s /* = {} */) {
+  if (!a_min.has_value() && !a_max.has_value()) {
+    throw std::invalid_argument("At most one of a_min and a_max may be None");
+  }
+  array result = astype(a, a.dtype(), s);
+  if (a_min.has_value()) {
+    result = maximum(result, a_min.value(), s);
+  }
+  if (a_max.has_value()) {
+    result = minimum(result, a_max.value(), s);
+  }
+  return result;
+}
+
 array concatenate(
     const std::vector<array>& arrays,
     int axis,
@@ -550,11 +678,11 @@ array concatenate(
     shape[ax] += a.shape(ax);
   }
 
+  // Promote all the arrays to the same type
+  auto dtype = result_type(arrays);
+
   return array(
-      shape,
-      arrays[0].dtype(),
-      std::make_unique<Concatenate>(to_stream(s), ax),
-      arrays);
+      shape, dtype, std::make_unique<Concatenate>(to_stream(s), ax), arrays);
 }
 
 array concatenate(
@@ -567,6 +695,29 @@ array concatenate(
   return concatenate(flat_inputs, 0, s);
 }
 
+/** Stack arrays along a new axis */
+array stack(
+    const std::vector<array>& arrays,
+    int axis,
+    StreamOrDevice s /* = {} */) {
+  if (arrays.empty()) {
+    throw std::invalid_argument("No arrays provided for stacking");
+  }
+  if (!is_same_shape(arrays)) {
+    throw std::invalid_argument("All arrays must have the same shape");
+  }
+  int normalized_axis = normalize_axis(axis, arrays[0].ndim() + 1);
+  std::vector<array> new_arrays;
+  new_arrays.reserve(arrays.size());
+  for (auto& a : arrays) {
+    new_arrays.emplace_back(expand_dims(a, normalized_axis, s));
+  }
+  return concatenate(new_arrays, axis, s);
+}
+array stack(const std::vector<array>& arrays, StreamOrDevice s /* = {} */) {
+  return stack(arrays, 0, s);
+}
+
 /** Pad an array with a constant value */
 array pad(
     const array& a,
@@ -653,6 +804,53 @@ array pad(
       s);
 }
 
+array moveaxis(
+    const array& a,
+    int source,
+    int destination,
+    StreamOrDevice s /* = {} */) {
+  auto check_ax = [&a](int ax) {
+    auto ndim = static_cast<int>(a.ndim());
+    if (ax < -ndim || ax >= ndim) {
+      std::ostringstream msg;
+      msg << "[moveaxis] Invalid axis " << ax << " for array with " << ndim
+          << " dimensions.";
+      throw std::out_of_range(msg.str());
+    }
+    return ax < 0 ? ax + ndim : ax;
+  };
+  source = check_ax(source);
+  destination = check_ax(destination);
+  std::vector<int> reorder(a.ndim());
+  std::iota(reorder.begin(), reorder.end(), 0);
+  reorder.erase(reorder.begin() + source);
+  reorder.insert(reorder.begin() + destination, source);
+  return transpose(a, reorder, s);
+}
+
+array swapaxes(
+    const array& a,
+    int axis1,
+    int axis2,
+    StreamOrDevice s /* = {} */) {
+  auto check_ax = [&a](int ax) {
+    auto ndim = static_cast<int>(a.ndim());
+    if (ax < -ndim || ax >= ndim) {
+      std::ostringstream msg;
+      msg << "[swapaxes] Invalid axis " << ax << " for array with " << ndim
+          << " dimensions.";
+      throw std::out_of_range(msg.str());
+    }
+    return ax < 0 ? ax + ndim : ax;
+  };
+  axis1 = check_ax(axis1);
+  axis2 = check_ax(axis2);
+  std::vector<int> reorder(a.ndim());
+  std::iota(reorder.begin(), reorder.end(), 0);
+  std::swap(reorder[axis1], reorder[axis2]);
+  return transpose(a, reorder, s);
+}
+
 array transpose(
     const array& a,
     std::vector<int> axes,
@@ -1438,6 +1636,34 @@ array operator/(const array& a, double b) {
   return divide(a, array(b));
 }
 
+array floor_divide(
+    const array& a,
+    const array& b,
+    StreamOrDevice s /* = {} */) {
+  auto dtype = promote_types(a.dtype(), b.dtype());
+  if (is_floating_point(dtype)) {
+    return floor(divide(a, b, s), s);
+  }
+
+  auto inputs = broadcast_arrays({astype(a, dtype, s), astype(b, dtype, s)}, s);
+  return array(
+      inputs[0].shape(), dtype, std::make_unique<Divide>(to_stream(s)), inputs);
+}
+
+array remainder(const array& a, const array& b, StreamOrDevice s /* = {} */) {
+  auto dtype = promote_types(a.dtype(), b.dtype());
+  auto inputs = broadcast_arrays(
+      {astype(a, dtype, s), astype(b, dtype, to_stream(s))}, s);
+  return array(
+      inputs[0].shape(),
+      dtype,
+      std::make_unique<Remainder>(to_stream(s)),
+      inputs);
+}
+array operator%(const array& a, const array& b) {
+  return remainder(a, b);
+}
+
 array maximum(const array& a, const array& b, StreamOrDevice s /* = {} */) {
   auto out_type = promote_types(a.dtype(), b.dtype());
   auto inputs =
@@ -1460,6 +1686,21 @@ array minimum(const array& a, const array& b, StreamOrDevice s /* = {} */) {
       inputs);
 }
 
+array floor(const array& a, StreamOrDevice s /* = {} */) {
+  if (a.dtype() == complex64) {
+    throw std::invalid_argument("[floor] Not supported for complex64.");
+  }
+  return array(
+      a.shape(), a.dtype(), std::make_unique<Floor>(to_stream(s)), {a});
+}
+
+array ceil(const array& a, StreamOrDevice s /* = {} */) {
+  if (a.dtype() == complex64) {
+    throw std::invalid_argument("[floor] Not supported for complex64.");
+  }
+  return array(a.shape(), a.dtype(), std::make_unique<Ceil>(to_stream(s)), {a});
+}
+
 array square(const array& a, StreamOrDevice s /* = {} */) {
   return array(
       a.shape(), a.dtype(), std::make_unique<Square>(to_stream(s)), {a});
@@ -1628,6 +1869,21 @@ array stop_gradient(const array& a, StreamOrDevice s /* = {} */) {
       a.shape(), a.dtype(), std::make_unique<StopGradient>(to_stream(s)), {a});
 }
 
+array round(const array& a, int decimals, StreamOrDevice s /* = {} */) {
+  if (decimals == 0) {
+    return array(
+        a.shape(), a.dtype(), std::make_unique<Round>(to_stream(s)), {a});
+  }
+
+  auto dtype = at_least_float(a.dtype());
+  float scale = std::pow(10, decimals);
+  auto result = multiply(a, array(scale, dtype), s);
+  result = round(result, 0, s);
+  result = multiply(result, array(1 / scale, dtype), s);
+
+  return astype(result, a.dtype(), s);
+}
+
 array matmul(
     const array& in_a,
     const array& in_b,
@@ -2322,4 +2578,184 @@ array conv2d(
       {in, wt});
 }
 
+array quantized_matmul(
+    const array& in_x,
+    const array& w,
+    const array& scales,
+    const array& biases,
+    int group_size /* = 64 */,
+    int bits /* = 4 */,
+    StreamOrDevice s /* = {} */) {
+  auto x = in_x;
+
+  if (w.dtype() != uint32) {
+    std::ostringstream msg;
+    msg << "[quantized_matmul] The weight matrix should be uint32 "
+        << "but received" << w.dtype();
+    throw std::invalid_argument(msg.str());
+  }
+  if (w.ndim() != 2) {
+    std::ostringstream msg;
+    msg << "[quantized_matmul] Batched quantized matmul is not supported for now "
+        << "received w with shape " << w.shape();
+    throw std::invalid_argument(msg.str());
+  }
+
+  // Keep x's batch dimensions to reshape it back after the matmul
+  auto original_shape = x.shape();
+  int x_inner_dims = original_shape.back();
+  original_shape.pop_back();
+
+  // Reshape x into a matrix if it isn't already one
+  if (x.ndim() != 2) {
+    x = reshape(x, {-1, x_inner_dims}, s);
+  }
+
+  int w_inner_dims = w.shape(0) * (32 / bits);
+  if (w_inner_dims != x_inner_dims) {
+    std::ostringstream msg;
+    msg << "[quantized_matmul] Last dimension of first input with "
+        << "shape (..., " << x_inner_dims
+        << ") does not match the expanded first "
+        << "dimension of the quantized matrix " << w_inner_dims
+        << ", computed from shape " << w.shape()
+        << " with group_size=" << group_size << " and bits=" << bits;
+    throw std::invalid_argument(msg.str());
+  }
+
+  int n_groups = x_inner_dims / group_size;
+  if (scales.shape(-1) != n_groups || biases.shape(-1) != n_groups) {
+    std::ostringstream msg;
+    msg << "[quantized_matmul] Scales and biases provided do not match the "
+        << "quantization arguments (group_size=" << group_size
+        << ", bits=" << bits << "). Expected shapes (" << w.shape(1) << ", "
+        << x_inner_dims / group_size
+        << "), but got scales.shape=" << scales.shape()
+        << " and biases.shape=" << biases.shape();
+    throw std::invalid_argument(msg.str());
+  }
+
+  auto out = array(
+      {x.shape(0), w.shape(1)},
+      x.dtype(),
+      std::make_unique<QuantizedMatmul>(to_stream(s), group_size, bits),
+      {x, w, scales, biases});
+
+  // If needed reshape x to the original batch shape
+  if (original_shape.size() != 1) {
+    original_shape.push_back(w.shape(1));
+    out = reshape(out, original_shape, s);
+  }
+
+  return out;
+}
+
+std::tuple<array, array, array> quantize(
+    const array& w,
+    int group_size /* = 64 */,
+    int bits /* = 4 */,
+    StreamOrDevice s /* = {} */) {
+  if (w.ndim() != 2) {
+    throw std::invalid_argument("[quantize] Only matrices supported for now");
+  }
+
+  if ((w.shape(0) % 32) != 0) {
+    throw std::invalid_argument(
+        "[quantize] All dimensions should be divisible by 32 for now");
+  }
+
+  if ((w.shape(-1) % group_size) != 0) {
+    std::ostringstream msg;
+    msg << "[quantize] The last dimension of the matrix needs to be divisible by "
+        << "the quantization group size " << group_size
+        << ". However the provided "
+        << " matrix has shape " << w.shape();
+    throw std::invalid_argument(msg.str());
+  }
+
+  // Compute some constants used for the quantization
+  int n_bins = (1 << bits) - 1; // 2**bits - 1
+  int el_per_int = 32 / bits;
+  array shifts = power(array(2, uint32), arange(0, 32, bits, uint32, s), s);
+  shifts = reshape(shifts, {1, 1, -1}, s);
+
+  // Compute scales and biases
+  array packed_w =
+      reshape(w, {w.shape(0), w.shape(1) / group_size, group_size}, s);
+  array w_max = max(packed_w, /* axis= */ -1, /* keepdims= */ true, s);
+  array w_min = min(packed_w, /* axis= */ -1, /* keepdims= */ true, s);
+  array delta = divide(subtract(w_max, w_min, s), array(n_bins, w.dtype()), s);
+  array scales = squeeze(delta, -1, s);
+  array biases = squeeze(w_min, -1, s);
+
+  // Quantize and pack w
+  packed_w =
+      astype(round(divide(subtract(packed_w, w_min, s), delta, s), s), uint32);
+  packed_w = reshape(packed_w, {w.shape(0), -1, el_per_int}, s);
+  packed_w = sum(
+      multiply(packed_w, shifts, s), /* axis= */ 2, /* keepdims= */ false, s);
+
+  return std::make_tuple(packed_w, scales, biases);
+}
+
+array dequantize(
+    const array& w,
+    const array& scales,
+    const array& biases,
+    int group_size /* = 64 */,
+    int bits /* = 4 */,
+    StreamOrDevice s /* = {} */) {
+  if (w.ndim() != 2 || scales.ndim() != 2 || biases.ndim() != 2) {
+    throw std::invalid_argument("[dequantize] Only matrices supported for now");
+  }
+
+  if ((w.shape(0) % 32) != 0) {
+    throw std::invalid_argument(
+        "[dequantize] All dimensions should be divisible by 32 for now");
+  }
+
+  if (w.shape(0) != scales.shape(0) || w.shape(0) != biases.shape(0)) {
+    throw std::invalid_argument(
+        "[dequantize] Shape of scales and biases does not match the matrix");
+  }
+
+  if (w.dtype() != uint32) {
+    throw std::invalid_argument(
+        "[dequantize] The matrix should be given as a uint32");
+  }
+
+  // Compute some constants for the dequantization
+  int el_per_int = 32 / bits;
+
+  if (w.shape(1) * el_per_int != scales.shape(1) * group_size) {
+    std::ostringstream msg;
+    msg << "[dequantize] Shape of scales and biases does not match the matrix "
+        << "given the quantization parameters. Provided matrix of shape "
+        << w.shape() << " and scales/biases of shape " << scales.shape()
+        << " with group_size=" << group_size << " and bits=" << bits << ".";
+    throw std::invalid_argument(msg.str());
+  }
+
+  // Extract the pieces from the passed quantized matrix
+  std::vector<array> parts;
+  for (int start = 0; start < 32; start += bits) {
+    // TODO: Implement bitwise operators for integral types
+    int shift_left = 32 - (start + bits);
+    int shift_right = shift_left + start;
+    array p = multiply(w, array(1 << shift_left, uint32), s);
+    p = floor_divide(p, array(1 << shift_right, uint32), s);
+    p = expand_dims(p, -1, s);
+    parts.push_back(p);
+  }
+  array w_full = concatenate(parts, -1, s);
+
+  // Dequantize
+  w_full = reshape(w_full, {w.shape(0), -1, group_size}, s);
+  w_full = multiply(w_full, expand_dims(scales, -1, s), s);
+  w_full = add(w_full, expand_dims(biases, -1, s), s);
+  w_full = reshape(w_full, {w.shape(0), -1}, s);
+
+  return w_full;
+}
+
 } // namespace mlx::core

mlx/ops.h

@@ -2,6 +2,7 @@
 
 #pragma once
 
+#include <optional>
 #include <variant>
 
 #include "array.h"
@@ -19,7 +20,7 @@ Stream to_stream(StreamOrDevice s);
 
 /**
  * A 1D array of numbers starting at `start` (optional),
- * stopping at stop, stepping by `step` (optional). **/
+ * stopping at stop, stepping by `step` (optional). */
 array arange(
     double start,
     double stop,
@@ -36,6 +37,14 @@ array arange(int start, int stop, int step, StreamOrDevice s = {});
 array arange(int start, int stop, StreamOrDevice s = {});
 array arange(int stop, StreamOrDevice s = {});
 
+/** A 1D array of `num` evenly spaced numbers in the range `[start, stop]` */
+array linspace(
+    double start,
+    double stop,
+    int num = 50,
+    Dtype dtype = float32,
+    StreamOrDevice s = {});
+
 /** Convert an array to the given data type. */
 array astype(const array& a, Dtype dtype, StreamOrDevice s = {});
 
@@ -87,11 +96,52 @@ inline array ones(const std::vector<int>& shape, StreamOrDevice s = {}) {
 }
 array ones_like(const array& a, StreamOrDevice s = {});
 
+/** Fill an array of the given shape (n,m) with ones in the specified diagonal
+ * k, and zeros everywhere else. */
+array eye(int n, int m, int k, Dtype dtype, StreamOrDevice s = {});
+inline array eye(int n, Dtype dtype, StreamOrDevice s = {}) {
+  return eye(n, n, 0, dtype, s);
+}
+inline array eye(int n, int m, StreamOrDevice s = {}) {
+  return eye(n, m, 0, float32, s);
+}
+inline array eye(int n, int m, int k, StreamOrDevice s = {}) {
+  return eye(n, m, k, float32, s);
+}
+inline array eye(int n, StreamOrDevice s = {}) {
+  return eye(n, n, 0, float32, s);
+}
+
+/** Create a square matrix of shape (n,n) of zeros, and ones in the major
+ * diagonal. */
+array identity(int n, Dtype dtype, StreamOrDevice s = {});
+inline array identity(int n, StreamOrDevice s = {}) {
+  return identity(n, float32, s);
+}
+
+array tri(int n, int m, int k, Dtype type, StreamOrDevice s = {});
+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 manipulation */
 
 /** Reshape an array to the given shape. */
 array reshape(const array& a, std::vector<int> shape, StreamOrDevice s = {});
 
+/** Flatten the dimensions in the range `[start_axis, end_axis]` . */
+array flatten(
+    const array& a,
+    int start_axis,
+    int end_axis = -1,
+    StreamOrDevice s = {});
+
+/** Flatten the array to 1D. */
+array flatten(const array& a, StreamOrDevice s = {});
+
 /** Remove singleton dimensions at the given axes. */
 array squeeze(
     const array& a,
@@ -144,6 +194,15 @@ std::vector<array> split(
 std::vector<array>
 split(const array& a, const std::vector<int>& indices, StreamOrDevice s = {});
 
+/**
+ * Clip (limit) the values in an array.
+ */
+array clip(
+    const array& a,
+    const std::optional<array>& a_min = std::nullopt,
+    const std::optional<array>& a_max = std::nullopt,
+    StreamOrDevice s = {});
+
 /** Concatenate arrays along a given axis. */
 array concatenate(
     const std::vector<array>& arrays,
@@ -151,6 +210,10 @@ array concatenate(
     StreamOrDevice s = {});
 array concatenate(const std::vector<array>& arrays, StreamOrDevice s = {});
 
+/** Stack arrays along a new axis. */
+array stack(const std::vector<array>& arrays, int axis, StreamOrDevice s = {});
+array stack(const std::vector<array>& arrays, StreamOrDevice s = {});
+
 /** Permutes the dimensions according to the given axes. */
 array transpose(const array& a, std::vector<int> axes, StreamOrDevice s = {});
 inline array transpose(
@@ -160,6 +223,16 @@ inline array transpose(
   return transpose(a, std::vector<int>(axes), s);
 }
 
+/** Swap two axes of an array. */
+array swapaxes(const array& a, int axis1, int axis2, StreamOrDevice s = {});
+
+/** Move an axis of an array. */
+array moveaxis(
+    const array& a,
+    int source,
+    int destination,
+    StreamOrDevice s = {});
+
 /** Pad an array with a constant value */
 array pad(
     const array& a,
@@ -636,12 +709,33 @@ array operator/(const array& a, const array& b);
 array operator/(double a, const array& b);
 array operator/(const array& a, double b);
 
+/** Compute integer division. Equivalent to doing floor(a / x). */
+array floor_divide(const array& a, const array& b, StreamOrDevice s = {});
+
+/** Compute the element-wise remainder of division */
+array remainder(const array& a, const array& b, StreamOrDevice s = {});
+array operator%(const array& a, const array& b);
+template <typename T>
+array operator%(T a, const array& b) {
+  return remainder(array(a), b);
+}
+template <typename T>
+array operator%(const array& a, T b) {
+  return remainder(a, array(b));
+}
+
 /** Element-wise maximum between two arrays. */
 array maximum(const array& a, const array& b, StreamOrDevice s = {});
 
 /** Element-wise minimum between two arrays. */
 array minimum(const array& a, const array& b, StreamOrDevice s = {});
 
+/** Floor the element of an array. **/
+array floor(const array& a, StreamOrDevice s = {});
+
+/** Ceil the element of an array. **/
+array ceil(const array& a, StreamOrDevice s = {});
+
 /** Square the elements of an array. */
 array square(const array& a, StreamOrDevice s = {});
 
@@ -711,6 +805,12 @@ array erfinv(const array& a, StreamOrDevice s = {});
 /** Stop the flow of gradients. */
 array stop_gradient(const array& a, StreamOrDevice s = {});
 
+/** Round a floating point number */
+array round(const array& a, int decimals, StreamOrDevice s = {});
+inline array round(const array& a, StreamOrDevice s = {}) {
+  return round(a, 0, s);
+}
+
 /** Matrix-matrix multiplication. */
 array matmul(const array& a, const array& b, StreamOrDevice s = {});
 
@@ -931,4 +1031,30 @@ 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,
+    const array& w,
+    const array& scales,
+    const array& biases,
+    int group_size = 64,
+    int bits = 4,
+    StreamOrDevice s = {});
+
+/** Quantize a matrix along its last axis */
+std::tuple<array, array, array> quantize(
+    const array& w,
+    int group_size = 64,
+    int bits = 4,
+    StreamOrDevice s = {});
+
+/** Dequantize a matrix produced by quantize() */
+array dequantize(
+    const array& w,
+    const array& scales,
+    const array& biases,
+    int group_size = 64,
+    int bits = 4,
+    StreamOrDevice s = {});
+
 } // namespace mlx::core

mlx/primitives.cpp

@@ -1,5 +1,4 @@
 // Copyright © 2023 Apple Inc.
-
 #include <algorithm>
 #include <cassert>
 #include <cmath>
@@ -441,6 +440,30 @@ bool Broadcast::is_equivalent(const Primitive& other) const {
   return shape_ == b_other.shape_;
 }
 
+std::vector<array> Ceil::vjp(
+    const std::vector<array>& primals,
+    const array& cotan,
+    const std::vector<int>& argnums) {
+  return {jvp(primals, {cotan}, argnums)};
+}
+
+array Ceil::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  assert(primals.size() == 1);
+  assert(argnums.size() == 1);
+  return zeros_like(primals[0], stream());
+}
+
+std::pair<array, int> Ceil::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  assert(inputs.size() == 1);
+  assert(axes.size() == 1);
+  return {ceil(inputs[0], stream()), axes[0]};
+}
+
 std::vector<array> Concatenate::vjp(
     const std::vector<array>& primals,
     const array& cotan,
@@ -488,7 +511,26 @@ array Concatenate::jvp(
 std::pair<array, int> Concatenate::vmap(
     const std::vector<array>& inputs,
     const std::vector<int>& axes) {
-  throw std::runtime_error("Concatenate vmap is NYI.");
+  std::vector<array> t_inputs;
+  // Find the first vmapped input
+  int i = 0;
+  for (; i < axes.size(); i++) {
+    t_inputs.push_back(inputs[i]);
+    if (axes[i] >= 0) {
+      break;
+    }
+  }
+  auto out_ax = axes[i++];
+  // Move vmap axes to the same spot.
+  for (; i < axes.size(); ++i) {
+    if (out_ax != axes[i] && axes[i] >= 0) {
+      t_inputs.push_back(moveaxis(inputs[i], axes[i], out_ax, stream()));
+    } else {
+      t_inputs.push_back(inputs[i]);
+    }
+  }
+  auto axis = axis_ + (axis_ >= out_ax);
+  return {concatenate(t_inputs, axis, stream()), out_ax};
 }
 
 bool Concatenate::is_equivalent(const Primitive& other) const {
@@ -738,6 +780,51 @@ std::pair<array, int> Divide::vmap(
   return {divide(a, b, stream()), to_ax};
 }
 
+std::vector<array> Remainder::vjp(
+    const std::vector<array>& primals,
+    const array& cotan,
+    const std::vector<int>& argnums) {
+  std::vector<array> vjps;
+  for (auto arg : argnums) {
+    if (arg == 0) {
+      vjps.push_back(cotan);
+    } else {
+      auto x_over_y = divide(primals[0], primals[1], stream());
+      x_over_y = floor(x_over_y, stream());
+      vjps.push_back(negative(multiply(x_over_y, cotan, stream()), stream()));
+    }
+  }
+  return vjps;
+}
+
+array Remainder::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  auto jvp_fun = [&](int i) {
+    int arg = argnums[i];
+    if (arg == 0) {
+      return tangents[i];
+    } else {
+      auto x_over_y = divide(primals[0], primals[1], stream());
+      x_over_y = floor(x_over_y, stream());
+      return negative(multiply(x_over_y, tangents[i], stream()), stream());
+    }
+  };
+  auto out = jvp_fun(0);
+  if (argnums.size() > 1) {
+    out = add(out, jvp_fun(1), stream());
+  }
+  return out;
+}
+
+std::pair<array, int> Remainder::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  auto [a, b, to_ax] = vmap_binary_op(inputs, axes, stream());
+  return {remainder(a, b, stream()), to_ax};
+}
+
 std::pair<array, int> Equal::vmap(
     const std::vector<array>& inputs,
     const std::vector<int>& axes) {
@@ -929,6 +1016,30 @@ array FFT::jvp(
   }
 }
 
+std::vector<array> Floor::vjp(
+    const std::vector<array>& primals,
+    const array& cotan,
+    const std::vector<int>& argnums) {
+  return {jvp(primals, {cotan}, argnums)};
+}
+
+array Floor::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  assert(primals.size() == 1);
+  assert(argnums.size() == 1);
+  return zeros_like(primals[0], stream());
+}
+
+std::pair<array, int> Floor::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  assert(inputs.size() == 1);
+  assert(axes.size() == 1);
+  return {floor(inputs[0], stream()), axes[0]};
+}
+
 std::vector<array> Full::vjp(
     const std::vector<array>& primals,
     const array& cotan,
@@ -961,7 +1072,53 @@ std::pair<array, int> Full::vmap(
 std::pair<array, int> Gather::vmap(
     const std::vector<array>& inputs,
     const std::vector<int>& axes) {
-  throw std::runtime_error("Gather vmap is NYI, please change slices instead");
+  auto& src = inputs[0];
+  std::vector<array> indices(inputs.begin() + 1, inputs.end());
+  auto gather_axes = axes_;
+  auto slice_sizes = slice_sizes_;
+  auto src_vmapped = axes[0] >= 0;
+  auto indices_vmapped =
+      std::any_of(axes.begin() + 1, axes.end(), [](int a) { return a >= 0; });
+  auto out_ax =
+      *std::find_if(axes.begin(), axes.end(), [](int a) { return a >= 0; });
+
+  // Reorder all the index arrays so the vmap axis is in the same spot.
+  for (int i = 1; i < axes.size(); ++i) {
+    if (out_ax != axes[i] && axes[i] >= 0) {
+      indices[i - 1] = moveaxis(indices[i - 1], axes[i], out_ax, stream());
+    }
+  }
+
+  if (src_vmapped) {
+    int max_dims = 0;
+    for (auto& idx : indices) {
+      max_dims = std::max(static_cast<int>(idx.ndim()), max_dims);
+    }
+    auto new_ax_loc =
+        std::find_if(gather_axes.begin(), gather_axes.end(), [&out_ax](int a) {
+          return a >= out_ax;
+        });
+    for (; new_ax_loc < gather_axes.end(); new_ax_loc++) {
+      (*new_ax_loc)++;
+    }
+    if (indices_vmapped) {
+      // Make a new index array for the vmapped dimension
+      // Reshape it so it broadcasts with other index arrays
+      // Update gather axes and slice sizes accordingly
+      auto shape = std::vector<int>(max_dims - out_ax, 1);
+      auto vmap_inds = arange(0, src.shape(out_ax), stream());
+      shape[0] = vmap_inds.shape(0);
+      vmap_inds = reshape(vmap_inds, shape, stream());
+      slice_sizes.insert(slice_sizes.begin() + out_ax, 1);
+      auto new_ax_idx = new_ax_loc - gather_axes.begin();
+      gather_axes.insert(new_ax_loc, out_ax);
+      indices.insert(indices.begin() + new_ax_idx, vmap_inds);
+    } else {
+      slice_sizes.insert(slice_sizes.begin() + axes[0], src.shape(axes[0]));
+      out_ax = max_dims + axes[0];
+    }
+  }
+  return {gather(src, indices, gather_axes, slice_sizes, stream()), out_ax};
 }
 
 std::vector<array> Gather::vjp(
@@ -1539,6 +1696,31 @@ std::pair<array, int> Power::vmap(
   return {power(a, b, stream()), to_ax};
 }
 
+std::pair<array, int> QuantizedMatmul::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  throw std::runtime_error("QuantizedMatmul::vmap NYI");
+}
+
+std::vector<array> QuantizedMatmul::vjp(
+    const std::vector<array>& primals,
+    const array& cotan,
+    const std::vector<int>& argnums) {
+  throw std::runtime_error("QuantizedMatmul::vjp NYI");
+}
+
+array QuantizedMatmul::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  throw std::runtime_error("QuantizedMatmul::vjp NYI");
+}
+
+bool QuantizedMatmul::is_equivalent(const Primitive& other) const {
+  const QuantizedMatmul& qm_other = static_cast<const QuantizedMatmul&>(other);
+  return group_size_ == qm_other.group_size_ && bits_ == qm_other.bits_;
+}
+
 std::pair<array, int> RandomBits::vmap(
     const std::vector<array>& inputs,
     const std::vector<int>& axes) {
@@ -1731,6 +1913,30 @@ bool Reduce::is_equivalent(const Primitive& other) const {
   return reduce_type_ == r_other.reduce_type_ && axes_ == r_other.axes_;
 }
 
+std::vector<array> Round::vjp(
+    const std::vector<array>& primals,
+    const array& cotan,
+    const std::vector<int>& argnums) {
+  return {jvp(primals, {cotan}, argnums)};
+}
+
+array Round::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  assert(primals.size() == 1);
+  assert(argnums.size() == 1);
+  return zeros_like(primals[0], stream());
+}
+
+std::pair<array, int> Round::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  assert(inputs.size() == 1);
+  assert(axes.size() == 1);
+  return {round(inputs[0], stream()), axes[0]};
+}
+
 std::pair<array, int> Scan::vmap(
     const std::vector<array>& inputs,
     const std::vector<int>& axes) {
@@ -1904,8 +2110,15 @@ std::pair<array, int> Sinh::vmap(
 std::pair<array, int> Slice::vmap(
     const std::vector<array>& inputs,
     const std::vector<int>& axes) {
-  // TODO implement
-  return {array(1.0f), axes[0]};
+  auto start = start_indices_;
+  auto stop = end_indices_;
+  auto strides = strides_;
+  auto ax = axes[0];
+  auto& input = inputs[0];
+  start.insert(start.begin() + ax, 0);
+  stop.insert(stop.begin() + ax, input.shape(ax));
+  strides.insert(strides.begin() + ax, 1);
+  return {slice(input, start, stop, strides, stream()), ax};
 }
 
 std::vector<array> Slice::vjp(

mlx/primitives.h

@@ -72,7 +72,7 @@ class Primitive {
       const std::vector<int>& argnums);
 
   /**
-   * The primitive must know how to vectorize itself accross
+   * The primitive must know how to vectorize itself across
    * the given axes. The output is a pair containing the array
    * representing the vectorized computation and the axis which
    * corresponds to the output vectorized dimension.
@@ -404,6 +404,25 @@ class Broadcast : public Primitive {
   void eval(const std::vector<array>& inputs, array& out);
 };
 
+class Ceil : public Primitive {
+ public:
+  explicit Ceil(Stream stream) : Primitive(stream){};
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  std::pair<array, int> vmap(
+      const std::vector<array>& inputs,
+      const std::vector<int>& axes) override;
+
+  DEFINE_GRADS()
+  DEFINE_PRINT(Ceil)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+
+ private:
+  void eval(const std::vector<array>& inputs, array& out);
+};
+
 class Concatenate : public Primitive {
  public:
   explicit Concatenate(Stream stream, int axis)
@@ -536,6 +555,25 @@ class Divide : public Primitive {
   void eval(const std::vector<array>& inputs, array& out);
 };
 
+class Remainder : public Primitive {
+ public:
+  explicit Remainder(Stream stream) : Primitive(stream){};
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  std::pair<array, int> vmap(
+      const std::vector<array>& inputs,
+      const std::vector<int>& axes) override;
+
+  DEFINE_GRADS()
+  DEFINE_PRINT(Remainder)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+
+ private:
+  void eval(const std::vector<array>& inputs, array& out);
+};
+
 class Equal : public Primitive {
  public:
   explicit Equal(Stream stream, bool equal_nan = false)
@@ -643,6 +681,25 @@ class FFT : public Primitive {
   void eval(const std::vector<array>& inputs, array& out);
 };
 
+class Floor : public Primitive {
+ public:
+  explicit Floor(Stream stream) : Primitive(stream){};
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  std::pair<array, int> vmap(
+      const std::vector<array>& inputs,
+      const std::vector<int>& axes) override;
+
+  DEFINE_GRADS()
+  DEFINE_PRINT(Floor)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+
+ private:
+  void eval(const std::vector<array>& inputs, array& out);
+};
+
 class Full : public Primitive {
  public:
   explicit Full(Stream stream) : Primitive(stream){};
@@ -1053,6 +1110,29 @@ class Power : public Primitive {
   void eval(const std::vector<array>& inputs, array& out);
 };
 
+class QuantizedMatmul : public Primitive {
+ public:
+  explicit QuantizedMatmul(Stream stream, int group_size, int bits)
+      : Primitive(stream), group_size_(group_size), bits_(bits){};
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  std::pair<array, int> vmap(
+      const std::vector<array>& inputs,
+      const std::vector<int>& axes) override;
+
+  DEFINE_GRADS()
+  DEFINE_PRINT(QuantizedMatmul)
+  bool is_equivalent(const Primitive& other) const override;
+
+ private:
+  int group_size_;
+  int bits_;
+
+  void eval(const std::vector<array>& inputs, array& out);
+};
+
 class RandomBits : public Primitive {
  public:
   explicit RandomBits(Stream stream, const std::vector<int>& shape, int width)
@@ -1149,6 +1229,25 @@ class Reduce : public Primitive {
   void eval(const std::vector<array>& inputs, array& out);
 };
 
+class Round : public Primitive {
+ public:
+  explicit Round(Stream stream) : Primitive(stream){};
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  std::pair<array, int> vmap(
+      const std::vector<array>& inputs,
+      const std::vector<int>& axes) override;
+
+  DEFINE_GRADS()
+  DEFINE_PRINT(Round)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+
+ private:
+  void eval(const std::vector<array>& inputs, array& out);
+};
+
 class Scan : public Primitive {
  public:
   enum ReduceType { Max, Min, Sum, Prod };

mlx/random.cpp

@@ -80,6 +80,16 @@ array split(const array& key, int num, StreamOrDevice s /* = {} */) {
   return bits({num, 2}, 4, key, s);
 }
 
+// Get the next representable value below 1.0 for half precision
+// floating point types (fp16, bf16)
+template <typename T>
+T below_one() {
+  T f = T(1.0);
+  uint16_t* m = (uint16_t*)&f;
+  *m -= 1;
+  return f;
+}
+
 array uniform(
     const array& low,
     const array& high,
@@ -87,9 +97,9 @@ array uniform(
     Dtype dtype /* = float32 */,
     const std::optional<array>& key /*= nullopt */,
     StreamOrDevice s /* = {} */) {
-  if (!is_floating_point(dtype)) {
+  if (!is_floating_point(dtype) && !is_complex(dtype)) {
     throw std::invalid_argument(
-        "Can only generate uniform numbers with floating point type.");
+        "Can only generate uniform numbers with real floating point type.");
   }
 
   auto stream = to_stream(s);
@@ -103,12 +113,29 @@ array uniform(
   }
   // Get random values between [0, nextafter(maxval, 0.0f)] since samples must
   // be in [low, high)
-  // TODO replace minimum with modulo uint32_t(nextafter(maxval, 0.0f)) to avoid
-  // clipping effects
-  float maxval = std::numeric_limits<uint32_t>::max();
-  auto upper = array(std::nextafter(maxval, 0.0f), dtype);
-  auto out = minimum(bits(shape, size_of(dtype), key, stream), upper, stream);
-  out = divide(out, array(maxval, dtype), stream);
+  auto get_limits = [&dtype]() {
+    switch (dtype) {
+      case float32:
+        return std::make_pair(
+            array(std::nextafter(1.0f, 0.0f), float32),
+            array(std::numeric_limits<uint32_t>::max(), float32));
+      case float16:
+        return std::make_pair(
+            array(below_one<float16_t>(), float16),
+            array(std::numeric_limits<uint16_t>::max(), float32));
+      case bfloat16:
+        return std::make_pair(
+            array(below_one<bfloat16_t>(), bfloat16),
+            array(std::numeric_limits<uint16_t>::max(), float32));
+      default:
+        throw std::runtime_error("[uniform] Unsupported type.");
+    }
+  };
+
+  auto [upper, maxval] = get_limits();
+  auto out = bits(shape, size_of(dtype), key, stream);
+  out = astype(divide(out, maxval, stream), dtype, stream);
+  out = minimum(out, upper, stream);
   return add(multiply(range, out, stream), low, stream);
 }
 

mlx/utils.cpp

@@ -49,6 +49,31 @@ std::vector<int> broadcast_shapes(
   return out_shape;
 }
 
+bool is_same_shape(const std::vector<array>& arrays) {
+  if (arrays.empty()) {
+    return true;
+  }
+  return std::all_of(arrays.begin() + 1, arrays.end(), [&](const array& a) {
+    return (a.shape() == arrays[0].shape());
+  });
+}
+
+int normalize_axis(int axis, int ndim) {
+  if (ndim <= 0) {
+    throw std::invalid_argument("Number of dimensions must be positive.");
+  }
+  if (axis < -ndim || axis >= ndim) {
+    std::ostringstream msg;
+    msg << "Axis " << axis << " is out of bounds for array with " << ndim
+        << " dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+  if (axis < 0) {
+    axis += ndim;
+  }
+  return axis;
+}
+
 std::ostream& operator<<(std::ostream& os, const Device& d) {
   os << "Device(";
   switch (d.type) {

mlx/utils.h

@@ -16,6 +16,15 @@ std::vector<int> broadcast_shapes(
     const std::vector<int>& s1,
     const std::vector<int>& s2);
 
+bool is_same_shape(const std::vector<array>& arrays);
+
+/**
+ * Returns the axis normalized to be in the range [0, ndim).
+ * Based on numpy's normalize_axis_index. See
+ * https://numpy.org/devdocs/reference/generated/numpy.lib.array_utils.normalize_axis_index.html
+ */
+int normalize_axis(int axis, int ndim);
+
 std::ostream& operator<<(std::ostream& os, const Device& d);
 std::ostream& operator<<(std::ostream& os, const Stream& s);
 std::ostream& operator<<(std::ostream& os, const Dtype& d);

python/mlx/extension.py

@@ -6,7 +6,7 @@ import subprocess
 import sys
 from pathlib import Path
 
-from setuptools import Extension, setup, find_namespace_packages
+from setuptools import Extension, find_namespace_packages, setup
 from setuptools.command.build_ext import build_ext
 
 import mlx

python/mlx/nn/__init__.py

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

python/mlx/nn/layers/__init__.py

@@ -1,23 +1,44 @@
 # Copyright © 2023 Apple Inc.
 
-from mlx.nn.layers.base import Module
 from mlx.nn.layers.activations import (
+    CELU,
+    ELU,
     GELU,
+    SELU,
+    LeakyReLU,
+    LogSigmoid,
+    Mish,
+    PReLU,
     ReLU,
+    ReLU6,
     SiLU,
+    Softplus,
+    Step,
+    celu,
+    elu,
     gelu,
     gelu_approx,
     gelu_fast_approx,
+    leaky_relu,
+    log_sigmoid,
+    mish,
+    prelu,
     relu,
+    relu6,
+    selu,
     silu,
+    softplus,
+    step,
 )
+from mlx.nn.layers.base import Module
 from mlx.nn.layers.containers import Sequential
 from mlx.nn.layers.convolution import Conv1d, Conv2d
 from mlx.nn.layers.dropout import Dropout
 from mlx.nn.layers.embedding import Embedding
 from mlx.nn.layers.linear import Linear
 from mlx.nn.layers.normalization import GroupNorm, LayerNorm, RMSNorm
-from mlx.nn.layers.positional_encoding import RoPE, SinusoidalPositionalEncoding
+from mlx.nn.layers.positional_encoding import ALiBi, RoPE, SinusoidalPositionalEncoding
+from mlx.nn.layers.quantized import QuantizedLinear
 from mlx.nn.layers.transformer import (
     MultiHeadAttention,
     TransformerEncoder,

python/mlx/nn/layers/activations.py

@@ -15,6 +15,15 @@ def _make_activation_module(f):
     return decorator
 
 
+def sigmoid(x):
+    r"""Applies the element-wise function:
+
+    .. math::
+        \text{Sigmoid}(x) = \sigma(x) = \frac{1}{1 + \exp(-x)}
+    """
+    return mx.sigmoid(x)
+
+
 def relu(x):
     """Applies the Rectified Linear Unit.
 
@@ -23,8 +32,49 @@ def relu(x):
     return mx.maximum(x, 0)
 
 
+def leaky_relu(x, negative_slope=0.01):
+    """Applies the Leaky Rectified Linear Unit.
+
+    Simply ``mx.maximum(negative_slope * x, x)``.
+    """
+    return mx.maximum(negative_slope * x, x)
+
+
+def elu(x, alpha=1.0):
+    """Applies the Exponential Linear Unit.
+
+    Simply ``mx.where(x > 0, x, alpha * (mx.exp(x) - 1))``.
+    """
+    return mx.where(x > 0, x, alpha * (mx.exp(x) - 1))
+
+
+def relu6(x):
+    r"""Applies the Rectified Linear Unit 6.
+
+    Applies :math:`\min(\max(x, 0), 6)` element wise.
+    """
+    return mx.minimum(mx.maximum(x, 0), 6.0)
+
+
+def softplus(x):
+    r"""Applies the Softplus function.
+
+    Applies :math:`\log(1 + \exp(x))` element wise.
+    """
+    return mx.logaddexp(x, 0)
+
+
+def celu(x, alpha=1.0):
+    r"""Applies the Continuously Differentiable Exponential Linear Unit.
+
+    Applies :math:`\max(0, x) + \min(0, \alpha * (\exp(x / \alpha) - 1))`
+    element wise.
+    """
+    return mx.maximum(x, 0.0) + alpha * (mx.exp(mx.minimum(x, 0.0) / alpha) - 1)
+
+
 def silu(x):
-    r"""Applies the Sigmoid Linear Unit.
+    r"""Applies the Sigmoid Linear Unit. Also known as Swish.
 
     Applies :math:`x \sigma(x)` element wise, where :math:`\sigma(\cdot)` is
     the logistic sigmoid.
@@ -32,8 +82,16 @@ def silu(x):
     return x * mx.sigmoid(x)
 
 
+def log_sigmoid(x):
+    r"""Applies the Log Sigmoid function.
+
+    Applies :math:`\log(\sigma(x)) = -\log(1 + e^{-x})` element wise.
+    """
+    return -softplus(-x)
+
+
 def gelu(x):
-    """Applies the Gaussian Error Linear Units function.
+    r"""Applies the Gaussian Error Linear Units function.
 
     .. math::
         \\textrm{GELU}(x) = x * \Phi(x)
@@ -80,16 +138,163 @@ def gelu_fast_approx(x):
     return x * mx.sigmoid(1.773 * x)
 
 
+@_make_activation_module
+class Sigmoid(Module):
+    pass
+
+
+def step(x: mx.array, threshold: float = 0.0):
+    r"""Applies the Step Activation Function.
+
+    This function implements a binary step activation, where the output is set
+    to 1 if the input is greater than a specified threshold, and 0 otherwise.
+
+    .. math::
+        \text{step}(x) = \begin{cases}
+        0 & \text{if } x < \text{threshold} \\
+        1 & \text{if } x \geq \text{threshold}
+        \end{cases}
+
+    Args:
+        threshold: The value to threshold at.
+    """
+
+    return mx.where(x > threshold, 1, 0)
+
+
+def selu(x):
+    r"""Applies the Scaled Exponential Linear Unit.
+
+    .. math::
+        \text{selu}(x) = \begin{cases}
+        \lambda x & \text{if } x > 0 \\
+        \lambda \alpha (\exp(x) - 1) & \text{if } x \leq 0
+        \end{cases}
+
+    where :math:`\lambda = 1.0507` and :math:`\alpha = 1.67326`.
+
+    See also :func:`elu`.
+    """
+    return elu(x, 1.67326) * 1.0507
+
+
+def prelu(x: mx.array, alpha: mx.array) -> mx.array:
+    r"""Applies the element-wise function:
+
+    .. math::
+        \text{PReLU}(x) = \max(0,x) + a * \min(0,x)
+
+    Here :math:`a` is an array.
+    """
+    return mx.maximum(0, x) + alpha * mx.minimum(0, x)
+
+
+def mish(x: mx.array) -> mx.array:
+    r"""Applies the Mish function, element-wise.
+    Mish: A Self Regularized Non-Monotonic Neural Activation Function.
+
+    Reference: https://arxiv.org/abs/1908.08681
+
+    .. math::
+        \text{Mish}(x) = x * \text{Tanh}(\text{Softplus}(x))
+
+    """
+    return x * mx.tanh(softplus(x))
+
+
+@_make_activation_module(mish)
+class Mish(Module):
+    pass
+
+
 @_make_activation_module(relu)
 class ReLU(Module):
     pass
 
 
+class LeakyReLU(Module):
+    r"""Applies the Leaky Rectified Linear Unit.
+
+    Simply ``mx.maximum(negative_slope * x, x)``.
+
+    Args:
+        negative_slope: Controls the angle of the negative slope. Default: 1e-2.
+    """
+
+    def __init__(self, negative_slope=1e-2):
+        super().__init__()
+        self._negative_slope = negative_slope
+
+    def __call__(self, x):
+        return leaky_relu(x, self._negative_slope)
+
+
+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.
+
+    Args:
+        alpha: the :math:`\alpha` value for the ELU formulation. Default: 1.0
+    """
+
+    def __init__(self, alpha=1.0):
+        super().__init__()
+        self._alpha = alpha
+
+    def __call__(self, x):
+        return elu(x, self._alpha)
+
+
+@_make_activation_module(relu6)
+class ReLU6(Module):
+    pass
+
+
+@_make_activation_module(softplus)
+class Softplus(Module):
+    pass
+
+
+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.
+
+    Args:
+        alpha: the :math:`\alpha` value for the CELU formulation. Default: 1.0
+    """
+
+    def __init__(self, alpha=1.0):
+        super().__init__()
+        self._alpha = alpha
+
+    def __call__(self, x):
+        return celu(x, self._alpha)
+
+
 @_make_activation_module(silu)
 class SiLU(Module):
     pass
 
 
+@_make_activation_module(log_sigmoid)
+class LogSigmoid(Module):
+    pass
+
+
+class PReLU(Module):
+    def __init__(self, num_parameters=1, init=0.25):
+        super().__init__()
+        self.weight = mx.full([num_parameters], init)
+
+    def __call__(self, x: mx.array):
+        return prelu(x, self.weight)
+
+
 class GELU(Module):
     r"""Applies the Gaussian Error Linear Units.
 
@@ -129,3 +334,45 @@ class GELU(Module):
 
     def __call__(self, x):
         return self._act(x)
+
+
+def tanh(x):
+    """Applies the hyperbolic tangent function.
+
+    Simply ``mx.tanh(x)``.
+    """
+    return mx.tanh(x)
+
+
+@_make_activation_module(tanh)
+class Tanh(Module):
+    pass
+
+
+class Step(Module):
+    r"""Applies the Step Activation Function.
+
+    This function implements a binary step activation, where the output is set
+    to 1 if the input is greater than a specified threshold, and 0 otherwise.
+
+    .. math::
+        \text{step}(x) = \begin{cases}
+        0 & \text{if } x < \text{threshold} \\
+        1 & \text{if } x \geq \text{threshold}
+        \end{cases}
+
+    Args:
+        threshold: The value to threshold at.
+    """
+
+    def __init__(self, threshold: float = 0.0):
+        super().__init__()
+        self.threshold = threshold
+
+    def __call__(self, x: mx.array):
+        return step(x, self.threshold)
+
+
+@_make_activation_module(selu)
+class SELU(Module):
+    pass

python/mlx/nn/layers/base.py

@@ -1,7 +1,7 @@
 # Copyright © 2023 Apple Inc.
 
 import textwrap
-from typing import Any, Callable, List, Union, Optional
+from typing import Any, Callable, List, Optional, Union
 
 import mlx.core as mx
 from mlx.utils import tree_flatten, tree_unflatten
@@ -258,6 +258,44 @@ class Module(dict):
         filter_fn = filter_fn or Module.valid_parameter_filter
         self.update(self.filter_and_map(filter_fn, map_fn))
 
+    def update_modules(self, modules: dict):
+        """Replace the child modules of this :class:`Module` instance with the
+        provided ones in the dict of dicts and lists.
+
+        It is the equivalent of :meth:`Module.update` but for modules instead
+        of parameters and allows us to flexibly edit complex architectures by
+        programmatically swapping layers.
+
+        The passed in parameters dictionary need not be a full dictionary
+        similar to :meth:`parameters`. Only the provided locations will be
+        updated.
+
+        Args:
+            modules (dict): A complete or partial dictionary of the modules
+                submodules.
+        """
+
+        def apply(dst, modules):
+            if isinstance(modules, dict):
+                for k in modules:
+                    if k in dst:
+                        current_value = dst[k]
+                        new_value = modules[k]
+                        if self.is_module(current_value) and self.is_module(new_value):
+                            dst[k] = new_value
+                        elif isinstance(current_value, (dict, list)):
+                            apply(current_value, new_value)
+            elif isinstance(modules, list):
+                for i in range(len(dst)):
+                    current_value = dst[i]
+                    new_value = modules[i]
+                    if self.is_module(current_value) and self.is_module(new_value):
+                        dst[i] = new_value
+                    elif isinstance(current_value, (dict, list)):
+                        apply(current_value, new_value)
+
+        apply(self, modules)
+
     def apply_to_modules(self, apply_fn: Callable[[str, "mlx.nn.Module"], Any]):
         """Apply a function to all the modules in this instance (including this
         instance).

python/mlx/nn/layers/convolution.py

@@ -78,7 +78,7 @@ class Conv2d(Module):
         out_channels (int): The number of output channels.
         kernel_size (int or tuple): The size of the convolution filters.
         stride (int or tuple, optional): The size of the stride when
-            applying the filter. Default: 0.
+            applying the filter. Default: 1.
         padding (int or tuple, optional): How many positions to 0-pad
             the input with. Default: 0.
         bias (bool, optional): If ``True`` add a learnable bias to the

python/mlx/nn/layers/linear.py

@@ -7,12 +7,21 @@ from mlx.nn.layers.base import Module
 
 
 class Linear(Module):
-    """Applies an affine transformation to the input.
+    r"""Applies an affine transformation to the input.
+
+    Concretely:
+
+    .. math::
+
+        y = W^\top x + b
+
+    where :math:`W` has shape ``[output_dims, input_dims]``.
 
     Args:
         input_dims (int): The dimensionality of the input features
         output_dims (int): The dimensionality of the output features
-        bias (bool): If set to False then the layer will not use a bias
+        bias (bool, optional): If set to ``False`` then the layer will
+          not use a bias. Default ``True``.
     """
 
     def __init__(self, input_dims: int, output_dims: int, bias: bool = True):

python/mlx/nn/layers/normalization.py

@@ -97,7 +97,7 @@ class GroupNorm(Module):
     where :math:`\gamma` and :math:`\beta` are learned per feature dimension
     parameters initialized at 1 and 0 respectively. However, the mean and
     variance are computed over the spatial dimensions and each group of
-    features. In particular, the input is split into num_groups accross the
+    features. In particular, the input is split into num_groups across the
     feature dimension.
 
     The feature dimension is assumed to be the last dimension and the dimensions

python/mlx/nn/layers/positional_encoding.py

@@ -18,15 +18,18 @@ class RoPE(Module):
 
     Args:
         dims (int): The feature dimensions to be rotated. If the input feature
-                    is larger than dims then the rest is left unchanged.
-        traditional (bool): If set to True choose the traditional
-                            implementation which is slightly less efficient.
+            is larger than dims then the rest is left unchanged.
+        traditional (bool, optional): If set to True choose the traditional
+            implementation which is slightly less efficient. Default: ``False``
+        base (float, optional): The base used to compute angular frequency for
+            each dimension in the positional encodings. Default: ``10000``
     """
 
-    def __init__(self, dims: int, traditional: bool = False):
+    def __init__(self, dims: int, traditional: bool = False, base: float = 10000):
         super().__init__()
         self.dims = dims
         self.traditional = traditional
+        self.base = base
 
     def _extra_repr(self):
         return f"{self.dims}, traditional={self.traditional}"
@@ -64,7 +67,7 @@ class RoPE(Module):
         x = mx.reshape(x, (-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
+            N, self.dims, offset=offset, base=self.base, dtype=x.dtype
         )
 
         rope = (
@@ -82,10 +85,7 @@ class RoPE(Module):
         positions = mx.arange(offset, N, dtype=dtype)
         freqs = mx.exp(-mx.arange(0.0, D, dtype=dtype) * (math.log(base) / D))
         theta = mx.reshape(positions, (-1, 1)) * mx.reshape(freqs, (1, -1))
-        costheta = mx.cos(theta)
-        sintheta = mx.sin(theta)
-
-        return costheta, sintheta
+        return mx.cos(theta), mx.sin(theta)
 
 
 class SinusoidalPositionalEncoding(Module):
@@ -142,3 +142,40 @@ class SinusoidalPositionalEncoding(Module):
             y = y * self.scale
 
         return y
+
+
+class ALiBi(Module):
+    @staticmethod
+    def create_alibi_matrix(
+        q_sequence_length: int,
+        k_sequence_length: int,
+        num_heads: int,
+        offset: int,
+        dtype=mx.float32,
+    ):
+        x1 = mx.arange(offset, q_sequence_length)
+        x2 = mx.arange(0, k_sequence_length)
+        distance_matrix = -mx.abs(
+            mx.expand_dims(x1[:, None] - x2[None, :], axis=(0, 1))
+        )
+        alibi_slope = ALiBi.create_alibi_slope(num_heads=num_heads)
+        alibi_mask = (distance_matrix * alibi_slope).astype(dtype)
+        return alibi_mask
+
+    @staticmethod
+    def create_alibi_slope(num_heads):
+        x = (2**8) ** (1 / num_heads)
+        out = mx.power(x, -mx.arange(1, num_heads + 1))
+        return mx.expand_dims(out, axis=(-1, -2))
+
+    def __call__(self, attention_scores, offset=0, mask=None):
+        alibi_mask = ALiBi.create_alibi_matrix(
+            q_sequence_length=attention_scores.shape[-2] + offset,
+            k_sequence_length=attention_scores.shape[-1],
+            num_heads=attention_scores.shape[1],
+            offset=offset,
+            dtype=attention_scores.dtype,
+        )
+        if mask is not None:
+            alibi_mask = alibi_mask + mask
+        return attention_scores + alibi_mask

python/mlx/nn/layers/quantized.py

@@ -0,0 +1,124 @@
+# Copyright © 2023 Apple Inc.
+
+import math
+
+import mlx.core as mx
+from mlx.nn.layers.base import Module
+from mlx.nn.layers.linear import Linear
+from mlx.utils import tree_flatten, tree_map
+
+
+class QuantizedLinear(Module):
+    """Applies an affine transformation to the input using a quantized weight matrix.
+
+    It is the quantized equivalent of :class:`mlx.nn.Linear`. For now its
+    parameters are frozen and will not be included in any gradient computation
+    but this will probably change in the future.
+
+    QuantizedLinear also provides two useful classmethods to convert linear
+    layers to QuantizedLinear layers.
+
+    - :meth:`from_linear` returns a QuantizedLinear layer that applies the same
+      linear transformation up to the quantization error.
+    - :meth:`quantize_module` swaps all the linear layers of the passed module
+      with QuantizedLinear ones.
+
+    Args:
+        input_dims (int): The dimensionality of the input features
+        output_dims (int): The dimensionality of the output features
+        bias (bool, optional): If set to ``False`` then the layer will not use
+            a bias. (default: True).
+        group_size (int, optional): The group size to use for the quantized
+            weight. See :func:`~mlx.core.quantize`. (default: 64)
+        bits (int, optional): The bit width to use for the quantized weight.
+            See :func:`~mlx.core.quantize`. (default: 4)
+    """
+
+    def __init__(
+        self,
+        input_dims: int,
+        output_dims: int,
+        bias: bool = True,
+        group_size: int = 64,
+        bits: int = 4,
+    ):
+        super().__init__()
+
+        # Quantization config
+        self.group_size = group_size
+        self.bits = bits
+
+        # Initialize the quantized weight
+        scale = math.sqrt(1 / input_dims)
+        weight = mx.random.uniform(
+            low=-scale,
+            high=scale,
+            shape=(output_dims, input_dims),
+        )
+        self.weight, self.scales, self.biases = mx.quantize(weight, group_size, bits)
+
+        # And bias if needed
+        if bias:
+            self.bias = mx.zeros((output_dims,))
+
+        # Freeze this model's parameters
+        self.freeze()
+
+    def unfreeze(self, *args, **kwargs):
+        """Wrap unfreeze so that we unfreeze any layers we might contain but
+        our parameters will remain frozen."""
+        super().unfreeze(*args, **kwargs)
+        self.freeze(recurse=False)
+
+    def _extra_repr(self):
+        out_dims, in_dims = self.weight.shape
+        in_dims *= 32 // self.bits
+        return (
+            f"input_dims={in_dims}, output_dims={out_dims}, bias={'bias' in self},"
+            f"group_size={self.group_size}, bits={self.bits}"
+        )
+
+    def __call__(self, x):
+        x = mx.quantized_matmul(
+            x,
+            self.weight.T,
+            scales=self.scales,
+            biases=self.biases,
+            group_size=self.group_size,
+            bits=self.bits,
+        )
+        if "bias" in self:
+            x = x + self.bias
+        return x
+
+    @classmethod
+    def from_linear(cls, linear_layer: Module, group_size: int = 64, bits: int = 4):
+        """Create a QuantizedLinear layer from the parameters of a provided
+        linear layer."""
+        output_dims, input_dims = linear_layer.weight.shape
+        ql = cls(input_dims, output_dims, False, group_size, bits)
+        ql.weight, ql.scales, ql.biases = mx.quantize(
+            linear_layer.weight, group_size, bits
+        )
+        if "bias" in linear_layer:
+            ql.bias = linear_layer.bias
+
+        return ql
+
+    @classmethod
+    def quantize_module(
+        cls,
+        model: Module,
+        group_size: int = 64,
+        bits: int = 4,
+        linear_class_predicate=lambda m: isinstance(m, Linear),
+    ):
+        def _quantize_if_linear(m):
+            if linear_class_predicate(m):
+                return cls.from_linear(m, group_size, bits)
+            else:
+                return m
+
+        leaves = model.leaf_modules()
+        leaves = tree_map(_quantize_if_linear, leaves, is_leaf=Module.is_module)
+        model.update_modules(leaves)

python/mlx/nn/layers/transformer.py

@@ -1,7 +1,7 @@
 # Copyright © 2023 Apple Inc.
 
 import math
-from typing import Optional
+from typing import Any, Optional
 
 import mlx.core as mx
 from mlx.nn.layers.base import Module
@@ -16,7 +16,7 @@ class MultiHeadAttention(Module):
     new values by aggregating information from the input values according to
     the similarities of the input queries and keys.
 
-    All inputs as well as the output are lineary projected without biases.
+    All inputs as well as the output are linearly projected without biases.
 
     MultiHeadAttention also expects an additive attention mask that should be
     broadcastable with (batch, num_heads, # queries, # keys). The mask should
@@ -43,12 +43,13 @@ class MultiHeadAttention(Module):
         value_input_dims: Optional[int] = None,
         value_dims: Optional[int] = None,
         value_output_dims: Optional[int] = None,
+        bias: bool = False,
     ):
         super().__init__()
 
         if (dims % num_heads) != 0:
             raise ValueError(
-                f"The input feature dimensions should be divisble by the number of heads ({dims} % {num_heads}) != 0"
+                f"The input feature dimensions should be divisible by the number of heads ({dims} % {num_heads}) != 0"
             )
 
         query_input_dims = query_input_dims or dims
@@ -58,10 +59,10 @@ class MultiHeadAttention(Module):
         value_output_dims = value_output_dims or dims
 
         self.num_heads = num_heads
-        self.query_proj = Linear(query_input_dims, dims, False)
-        self.key_proj = Linear(key_input_dims, dims, False)
-        self.value_proj = Linear(value_input_dims, value_dims, False)
-        self.out_proj = Linear(value_dims, value_output_dims, False)
+        self.query_proj = Linear(query_input_dims, dims, bias=bias)
+        self.key_proj = Linear(key_input_dims, dims, bias=bias)
+        self.value_proj = Linear(value_input_dims, value_dims, bias=bias)
+        self.out_proj = Linear(value_dims, value_output_dims, bias=bias)
 
     def __call__(self, queries, keys, values, mask=None):
         queries = self.query_proj(queries)
@@ -136,3 +137,85 @@ class TransformerEncoder(Module):
         x = self.ln(x)
 
         return x
+
+
+class TransformerDecoderLayer(Module):
+    def __init__(self, dims: int, num_heads: int, mlp_dims: Optional[int] = None):
+        super().__init__()
+        mlp_dims = mlp_dims or dims * 4
+        self.self_attention = MultiHeadAttention(dims, num_heads)
+        self.cross_attention = MultiHeadAttention(dims, num_heads)
+        self.ln1 = LayerNorm(dims)
+        self.ln2 = LayerNorm(dims)
+        self.ln3 = LayerNorm(dims)
+        self.linear1 = Linear(dims, mlp_dims)
+        self.linear2 = Linear(mlp_dims, dims)
+
+    def __call__(self, x, memory, x_mask, memory_mask):
+        y = self.ln1(x)
+        y = self.self_attention(y, y, y, x_mask)
+        x = x + y
+
+        y = self.ln2(x)
+        y = self.cross_attention(y, memory, memory, memory_mask)
+        x = x + y
+
+        y = self.ln3(x)
+        y = self.linear1(y)
+        y = mx.maximum(y, 0)
+        y = self.linear2(y)
+        x = x + y
+
+        return x
+
+
+class TransformerDecoder(Module):
+    def __init__(
+        self, num_layers: int, dims: int, num_heads: int, mlp_dims: Optional[int] = None
+    ):
+        super().__init__()
+        self.layers = [
+            TransformerDecoderLayer(dims, num_heads, mlp_dims)
+            for i in range(num_layers)
+        ]
+        self.ln = LayerNorm(dims)
+
+    def __call__(self, x, memory, x_mask, memory_mask):
+        for l in self.layers:
+            x = l(x, memory, x_mask, memory_mask)
+        x = self.ln(x)
+
+        return x
+
+
+class Transformer(Module):
+    def __init__(
+        self,
+        dims: int = 512,
+        num_heads: int = 8,
+        num_encoder_layers: int = 6,
+        num_decoder_layers: int = 6,
+        mlp_dims: Optional[int] = None,
+        custom_encoder: Optional[Any] = None,
+        custom_decoder: Optional[Any] = None,
+    ):
+        super().__init__()
+        if custom_encoder is not None:
+            self.encoder = custom_encoder
+        else:
+            self.encoder = TransformerEncoder(
+                num_encoder_layers, dims, num_heads, mlp_dims
+            )
+
+        if custom_decoder is not None:
+            self.decoder = custom_decoder
+        else:
+            self.decoder = TransformerDecoder(
+                num_decoder_layers, dims, num_heads, mlp_dims
+            )
+
+    def __call__(self, src, tgt, src_mask, tgt_mask, memory_mask):
+        memory = self.encoder(src, src_mask)
+        output = self.decoder(tgt, memory, tgt_mask, memory_mask)
+
+        return output

python/mlx/nn/losses.py

@@ -1,8 +1,285 @@
 # Copyright © 2023 Apple Inc.
 
 import mlx.core as mx
+from mlx.nn.layers.base import Module
 
 
-def cross_entropy(logits: mx.array, targets: mx.array, axis: int = -1):
+def cross_entropy(
+    logits: mx.array,
+    targets: mx.array,
+    weights: mx.array = None,
+    axis: int = -1,
+    label_smoothing: float = 0.0,
+    reduction: str = "none",
+) -> mx.array:
+    """
+    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``.
+        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:
+            ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'none'``.
+
+    Returns:
+        array: The computed cross entropy loss.
+    """
+    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)
-    return mx.logsumexp(logits, axis=axis) - score
+    logsumexp_logits = mx.logsumexp(logits, axis=axis)
+    if label_smoothing > 0:
+        # Adjust the true class score with label smoothing
+        adjusted_score = (1 - label_smoothing) * score
+
+        # Calculate the mean logit across the classes for smoothed loss
+        mean_logits = logits.mean(axis=axis)
+        smoothed_loss = -mean_logits * label_smoothing
+
+        # Combine the adjusted score and smoothed loss with the logsumexp logits
+        loss = logsumexp_logits - adjusted_score + smoothed_loss
+    else:
+        loss = logsumexp_logits - score
+
+    # Apply weights if provided
+    if weights is not None:
+        if weights.shape != targets.shape:
+            raise ValueError(
+                f"Weights with shape {weights.shape} is not the same as "
+                f"targets with shape {targets.shape}."
+            )
+        loss *= weights
+
+    # Apply reduction
+    return _reduce(loss, reduction)
+
+
+def binary_cross_entropy(
+    logits: mx.array, targets: mx.array, reduction: str = "none"
+) -> mx.array:
+    """
+    Computes the binary cross entropy loss.
+
+    Args:
+        logits (array): The unnormalized (pre-sigmoid) predicted logits.
+        targets (array): The binary target values in {0, 1}.
+        reduction (str, optional): Specifies the reduction to apply to the output:
+          ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'none'``.
+
+    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])
+        >>> targets = mx.array([0, 0, 1, 1])
+        >>> loss = nn.losses.binary_cross_entropy(inputs, targets, "mean")
+        >>> loss
+        array([0.612192], dtype=float32)
+    """
+    loss = mx.logaddexp(0.0, logits) - targets * logits
+    return _reduce(loss, reduction)
+
+
+def l1_loss(
+    predictions: mx.array, targets: mx.array, reduction: str = "mean"
+) -> mx.array:
+    """
+    Computes the L1 loss.
+
+    Args:
+        predictions (array): The predicted values.
+        targets (array): The target values.
+        reduction (str, optional): Specifies the reduction to apply to the output:
+          ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'mean'``.
+
+    Returns:
+        array: The computed L1 loss.
+    """
+    if predictions.shape != targets.shape:
+        raise ValueError(
+            f"Predictions shape {predictions.shape} does not match "
+            f"targets shape {targets.shape}."
+        )
+    loss = mx.abs(predictions - targets)
+
+    return _reduce(loss, reduction)
+
+
+def mse_loss(
+    predictions: mx.array, targets: mx.array, reduction: str = "mean"
+) -> mx.array:
+    """
+    Computes the mean squared error loss.
+
+    Args:
+        predictions (array): The predicted values.
+        targets (array): The target values.
+        reduction (str, optional): Specifies the reduction to apply to the output:
+          ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'mean'``.
+
+    Returns:
+        array: The computed mean squared error loss.
+    """
+    if predictions.shape != targets.shape:
+        raise ValueError(
+            f"Predictions shape {predictions.shape} does not match "
+            f"targets shape {targets.shape}."
+        )
+
+    assert (
+        predictions.shape == targets.shape
+    ), f"Shape of predictions {predictions.shape} and targets {targets.shape} must match"
+
+    loss = mx.square(predictions - targets)
+    return _reduce(loss, reduction)
+
+
+def nll_loss(
+    inputs: mx.array, targets: mx.array, axis: int = -1, reduction: str = "none"
+) -> mx.array:
+    """
+    Computes the negative log likelihood loss.
+
+    Args:
+        inputs (array): The predicted distribution in log space.
+        targets (array): The target values.
+        axis (int, optional): The distribution axis. Default: ``-1``.
+        reduction (str, optional): Specifies the reduction to apply to the output:
+          ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'none'``.
+
+    Returns:
+        array: The computed NLL loss.
+    """
+    loss = -mx.take_along_axis(inputs, targets[..., None], axis).squeeze(-1)
+
+    return _reduce(loss, reduction)
+
+
+def kl_div_loss(
+    inputs: mx.array, targets: mx.array, axis: int = -1, reduction: str = "none"
+) -> mx.array:
+    """
+    Computes the Kullback-Leibler divergence loss.
+
+    Computes the following when ``reduction == 'none'``:
+
+    .. code-block:: python
+
+        mx.exp(targets) * (targets - inputs).sum(axis)
+
+    Args:
+        inputs (array): Log probabilities for the predicted distribution.
+        targets (array): Log probabilities for the target distribution.
+        axis (int, optional): The distribution axis. Default: ``-1``.
+        reduction (str, optional): Specifies the reduction to apply to the output:
+          ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'none'``.
+
+    Returns:
+        array: The computed Kullback-Leibler divergence loss.
+    """
+    loss = mx.sum(mx.exp(targets) * (targets - inputs), axis)
+
+    return _reduce(loss, reduction)
+
+
+def smooth_l1_loss(
+    predictions: mx.array, targets: mx.array, beta: float = 1.0, reduction: str = "mean"
+) -> mx.array:
+    r"""
+    Computes the smooth L1 loss.
+
+    The smooth L1 loss is a variant of the L1 loss which replaces the absolute
+    difference with a squared difference when the absolute difference is less
+    than ``beta``.
+
+    The formula for the smooth L1 Loss is:
+
+    .. math::
+
+       l =
+          \begin{cases}
+            0.5 (x - y)^2, & \text{ if } & (x - y) < \beta \\
+            |x - y| - 0.5 \beta, &  & \text{otherwise}
+          \end{cases}
+
+    Args:
+        predictions (array): Predicted values.
+        targets (array): Ground truth values.
+        beta (float, optional): The threshold after which the loss changes
+          from the squared to the absolute difference. Default: ``1.0``.
+        reduction (str, optional): Specifies the reduction to apply to the output:
+          ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'mean'``.
+
+    Returns:
+        array: The computed smooth L1 loss.
+    """
+    if predictions.shape != targets.shape:
+        raise ValueError(
+            f"Predictions shape {predictions.shape} does not match "
+            f"targets shape {targets.shape}."
+        )
+
+    diff = predictions - targets
+    loss = mx.where(
+        diff < beta, 0.5 * mx.square(diff) / beta, mx.abs(diff) - 0.5 * beta
+    )
+
+    return _reduce(loss, reduction)
+
+
+def triplet_loss(
+    anchors: mx.array,
+    positives: mx.array,
+    negatives: mx.array,
+    axis: int = -1,
+    p: int = 2,
+    margin: float = 1.0,
+    eps: float = 1e-6,
+    reduction: str = "none",
+) -> mx.array:
+    r"""
+    Computes the triplet loss for a set of anchor, positive, and negative samples.
+    Margin is represented with alpha in the math section.
+
+    .. math::
+
+       L_{\text{triplet}} = \max\left(\|A - P\|_p - \|A - N\|_p + \alpha, 0\right)
+
+    Args:
+        anchors (array): The anchor samples.
+        positives (array): The positive samples.
+        negatives (array): The negative samples.
+        axis (int, optional): The distribution axis. Default: ``-1``.
+        p (int, optional): The norm degree for pairwise distance. Default: ``2``.
+        margin (float, optional): Margin for the triplet loss. Defaults to ``1.0``.
+        eps (float, optional): Small positive constant to prevent numerical instability. Defaults to ``1e-6``.
+        reduction (str, optional): Specifies the reduction to apply to the output:
+          ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'none'``.
+
+    Returns:
+        array: Computed triplet loss. If reduction is "none", returns a tensor of the same shape as input;
+                  if reduction is "mean" or "sum", returns a scalar tensor.
+    """
+    loss = mx.maximum(
+        mx.sqrt(mx.power(anchors - positives, p).sum(axis) + eps)
+        - mx.sqrt(mx.power(anchors - negatives, p).sum(axis) + eps)
+        + margin,
+        0,
+    )
+    return _reduce(loss, reduction)
+
+
+def _reduce(loss: mx.array, reduction: str = "none"):
+    if reduction == "mean":
+        return mx.mean(loss)
+    elif reduction == "sum":
+        return mx.sum(loss)
+    elif reduction == "none":
+        return loss
+    else:
+        raise ValueError("Invalid reduction. Must be 'none', 'mean', or 'sum'.")

python/mlx/optimizers.py

@@ -82,19 +82,36 @@ class SGD(Optimizer):
 
     .. math::
 
-        v_{t+1} &= \mu v_t + (1 - \mu) g_t \\
+        v_{t+1} &= \mu v_t + (1 - \tau) g_t \\
         w_{t+1} &= w_t - \lambda v_{t+1}
 
     Args:
-        learning_rate (float): The learning :math:`\lambda` for the update
-        momentum (float): The momentum strength :math:`\mu`
+        learning_rate (float): The learning rate :math:`\lambda`.
+        momentum (float, optional): The momentum strength :math:`\mu`. Default: ``0``
+        weight_decay (float, optional): The weight decay (L2 penalty). Default: ``0``
+        dampening (float, optional): Dampening for momentum :math:`\tau`. Default: ``0``
+        nesterov (bool, optional): Enables Nesterov momentum. Default: ``False``
     """
 
-    def __init__(self, learning_rate: float, momentum: float = 0.0):
+    def __init__(
+        self,
+        learning_rate: float,
+        momentum: float = 0.0,
+        weight_decay: float = 0.0,
+        dampening: float = 0.0,
+        nesterov: bool = False,
+    ):
+        if nesterov and (momentum <= 0 or dampening != 0):
+            raise ValueError(
+                "Nesterov momentum requires a momentum and zero dampening."
+            )
         super().__init__()
 
         self.learning_rate = learning_rate
         self.momentum = momentum
+        self.weight_decay = weight_decay
+        self.dampening = dampening
+        self.nesterov = nesterov
 
     def apply_single(
         self, gradient: mx.array, parameter: mx.array, state: OptimizerState
@@ -105,9 +122,175 @@ class SGD(Optimizer):
             return parameter - self.learning_rate * gradient
 
         v = state.get("v", mx.zeros_like(gradient))
-        v = self.momentum * v + (1 - self.momentum) * gradient
+
+        if self.weight_decay != 0:
+            gradient += self.weight_decay * parameter
+
+        v = self.momentum * v
+        if self.dampening > 0:
+            v += (1 - self.dampening) * gradient
+        else:
+            v += gradient
+
+        if self.nesterov:
+            update = gradient + self.momentum * v
+        else:
+            update = v
         state["v"] = v
-        return parameter - self.learning_rate * v
+        return parameter - self.learning_rate * update
+
+
+class RMSprop(Optimizer):
+    r"""Implementation of the RMSprop optimizer [1].
+
+    [1]: Tieleman, T. and Hinton, G. 2012. Lecture 6.5-rmsprop, coursera: Neural networks for machine learning
+
+    .. math::
+
+        v_{t+1} &= \alpha v_t + (1 - \alpha) g_t^2 \\
+        w_{t+1} &= w_t - \lambda \frac{g_t}{\sqrt{v_{t+1}} + \epsilon}
+
+    Args:
+        learning_rate (float): The learning rate :math:`\lambda`.
+        alpha (float, optional): The smoothing constant :math:`\alpha`.
+          Default: ``0.99``
+        eps (float, optional): The term :math:`\epsilon` added to the denominator
+          to improve numerical stability. Default: ``1e-8``
+    """
+
+    def __init__(self, learning_rate: float, alpha: float = 0.99, eps: float = 1e-8):
+        super().__init__()
+
+        self.learning_rate = learning_rate
+        self.alpha = alpha
+        self.eps = eps
+
+        if self.alpha < 0.0:
+            raise ValueError(
+                f"RMSprop alpha should be >=0, {self.alpha} was provided instead"
+            )
+        if self.eps < 0.0:
+            raise ValueError(
+                f"RMSprop epsilon should be >0, {self.eps} was provided instead"
+            )
+
+    def apply_single(
+        self, gradient: mx.array, parameter: mx.array, state: OptimizerState
+    ):
+        """Performs the RMSprop parameter update and stores :math:`v` in the optimizer state."""
+        lr = self.learning_rate
+        alpha = self.alpha
+        eps = self.eps
+
+        v = state.get("v", mx.zeros_like(gradient))
+        v = alpha * v + (1 - alpha) * mx.square(gradient)
+        state["v"] = v
+
+        return parameter - lr * gradient / (mx.sqrt(v) + eps)
+
+
+class Adagrad(Optimizer):
+    r"""Implementation of the Adagrad optimizer [1].
+
+    Our Adagrad implementation follows the original paper. In detail,
+
+    [1]: Duchi, J., Hazan, E. and Singer, Y., 2011. Adaptive subgradient methods
+    for online learning and stochastic optimization. JMLR 2011.
+
+    .. math::
+
+        v_{t+1} &= v_t + g_t^2 \\
+        w_{t+1} &= w_t - \lambda \frac{g_t}{\sqrt{v_{t+1}} + \epsilon}
+
+    Args:
+        learning_rate (float): The learning rate :math:`\lambda`.
+        eps (float, optional): The term :math:`\epsilon` added to the
+          denominator to improve numerical stability. Default: ``1e-8``
+    """
+
+    def __init__(self, learning_rate: float, eps: float = 1e-8):
+        super().__init__()
+
+        self.learning_rate = learning_rate
+        self.eps = eps
+
+        if self.eps < 0.0:
+            raise ValueError(
+                f"Adagrad epsilon should be >0, {self.eps} was provided instead"
+            )
+
+    def apply_single(
+        self, gradient: mx.array, parameter: mx.array, state: OptimizerState
+    ):
+        """Performs the Adagrad parameter update and stores :math:`v` in the
+        optimizer state."""
+        lr = self.learning_rate
+        eps = self.eps
+
+        v = state.get("v", mx.zeros_like(gradient))
+        v = v + mx.square(gradient)
+        state["v"] = v
+
+        return parameter - lr * gradient / (mx.sqrt(v) + eps)
+
+
+class AdaDelta(Optimizer):
+    r"""Implementation of the AdaDelta optimizer with learning rate[1].
+
+    Our AdaDelta implementation follows the original paper. In detail,
+
+    [1]: Zeiler, M.D., 2012. ADADELTA: an adaptive learning rate method. arXiv preprint arXiv:1212.5701.
+
+    .. math::
+
+        v_{t+1} &= \rho v_t + (1 - \rho) g_t^2 \\
+        \Delta w_{t+1} &= \frac{\sqrt{u_t + \epsilon}}{\sqrt{v_{t+1} + \epsilon}} g_t \\
+        u_{t+1} &= \rho u_t + (1 - \rho) \Delta w_{t+1}^2 \\
+        w_{t+1} &= w_t - \lambda \Delta w_{t+1}
+
+    Args:
+        learning_rate (float): The learning rate :math:`\lambda`.
+        rho (float, optional): The coefficient :math:`\rho` used for computing a
+            running average of squared gradients. Default: ``0.9``
+        eps (float, optional): The term :math:`\epsilon` added to the denominator to improve
+          numerical stability. Ddefault: `1e-8`
+    """
+
+    def __init__(self, learning_rate: float, rho: float = 0.9, eps: float = 1e-6):
+        super().__init__()
+
+        self.learning_rate = learning_rate
+        self.rho = rho
+        self.eps = eps
+        if self.rho < 0.0:
+            raise ValueError(
+                f"AdaDelta rho should be >=0, {self.rho} was provided instead"
+            )
+        if self.eps < 0.0:
+            raise ValueError(
+                f"AdaDelta epsilon should be >0, {self.eps} was provided instead"
+            )
+
+    def apply_single(
+        self, gradient: mx.array, parameter: mx.array, state: OptimizerState
+    ):
+        """Performs the AdaDelta parameter update and stores :math:`v` and
+        :math:`u` in the optimizer state."""
+        lr = self.learning_rate
+        rho = self.rho
+        eps = self.eps
+
+        v = state.get("v", mx.zeros_like(gradient))
+        u = state.get("s", mx.zeros_like(gradient))
+
+        v = rho * v + (1 - rho) * mx.square(gradient)
+        d = mx.sqrt(u + eps) / mx.sqrt(v + eps) * gradient
+        u = rho * u + (1 - rho) * mx.square(d)
+
+        state["v"] = v
+        state["u"] = u
+
+        return parameter - lr * d
 
 
 class Adam(Optimizer):
@@ -116,14 +299,22 @@ class Adam(Optimizer):
     Our Adam implementation follows the original paper and omits the bias
     correction in the first and second moment estimates. In detail,
 
+    [1]: Kingma, D.P. and Ba, J., 2015. Adam: A method for stochastic
+    optimization. ICLR 2015.
+
     .. math::
 
         m_{t+1} &= \beta_1 m_t + (1 - \beta_1) g_t \\
         v_{t+1} &= \beta_2 v_t + (1 - \beta_2) g_t^2 \\
         w_{t+1} &= w_t - \lambda \frac{m_{t+1}}{\sqrt{v_{t+1} + \epsilon}}
 
-    [1]: Kingma, D.P. and Ba, J., 2015. Adam: A method for stochastic
-    optimization. ICLR 2015.
+    Args:
+        learning_rate (float): The learning rate :math:`\lambda`.
+        betas (Tuple[float, float], optional): The coefficients
+          :math:`(\beta_1, \beta_2)` used for computing running averages of the
+          gradient and its square. Default: ``(0.9, 0.999)``
+        eps (float, optional): The term :math:`\epsilon` added to the
+          denominator to improve numerical stability. Default: ``1e-8``
     """
 
     def __init__(
@@ -152,3 +343,158 @@ class Adam(Optimizer):
         state["v"] = v
 
         return parameter - lr * m / (mx.sqrt(v) + eps)
+
+
+class AdamW(Adam):
+    r"""Implementation of 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
+    with a weight_decay (:math:`\lambda`) value:
+
+    [1]: Loshchilov, I. and Hutter, F., 2019. Decoupled weight decay
+    regularization. ICLR 2019.
+
+    .. math::
+
+        m_{t+1} &= \beta_1 m_t + (1 - \beta_1) g_t \\
+        v_{t+1} &= \beta_2 v_t + (1 - \beta_2) g_t^2 \\
+        w_{t+1} &= w_t - \alpha (\frac{m_{t+1}}{\sqrt{v_{t+1} + \epsilon}} + \lambda w_t)
+
+    Args:
+        learning_rate (float): The learning rate :math:`\alpha`.
+        betas (Tuple[float, float], optional): The coefficients
+          :math:`(\beta_1, \beta_2)` used for computing running averages of the
+          gradient and its square. Default: ``(0.9, 0.999)``
+        eps (float, optional): The term :math:`\epsilon` added to the
+          denominator to improve numerical stability. Default: ``1e-8``
+        weight_decay (float, optional): The weight decay :math:`\lambda`.
+          Default: ``0``.
+    """
+
+    def __init__(
+        self,
+        learning_rate: float,
+        betas: List[float] = [0.9, 0.999],
+        eps: float = 1e-8,
+        weight_decay: float = 0.01,
+    ):
+        super().__init__(learning_rate=learning_rate, betas=betas, eps=eps)
+        self.weight_decay = weight_decay
+
+    def apply_single(
+        self, gradient: mx.array, parameter: mx.array, state: OptimizerState
+    ):
+        """Performs the AdamW parameter update by modifying the parameters
+        passed into Adam.
+        """
+
+        return super().apply_single(
+            gradient, parameter * (1 - self.learning_rate * self.weight_decay), state
+        )
+
+
+class Adamax(Adam):
+    r"""Implementation of the Adamax optimizer. It is 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,
+
+    [1]: Kingma, D.P. and Ba, J., 2015. Adam: A method for stochastic
+    optimization. ICLR 2015.
+
+    .. math::
+
+        m_{t+1} &= \beta_1 m_t + (1 - \beta_1) g_t \\
+        v_{t+1} &= \max(\beta_2 v_t, |g_t|) \\
+        w_{t+1} &= w_t - \lambda \frac{m_{t+1}}{v_{t+1} + \epsilon}
+
+    Args:
+        learning_rate (float): The learning rate :math:`\lambda`.
+        betas (Tuple[float, float], optional): The coefficients
+          :math:`(\beta_1, \beta_2)` used for computing running averages of the
+          gradient and its square. Default: ``(0.9, 0.999)``
+        eps (float, optional): The term :math:`\epsilon` added to the
+          denominator to improve numerical stability. Default: ``1e-8``
+    """
+
+    def __init__(
+        self, learning_rate: float, betas: List[float] = [0.9, 0.999], eps: float = 1e-8
+    ):
+        super().__init__(learning_rate, betas, eps)
+
+    def apply_single(
+        self, gradient: mx.array, parameter: mx.array, state: OptimizerState
+    ):
+        """Performs the Adamax parameter update and stores :math:`v` and
+        :math:`m` in the optimizer state."""
+        lr = self.learning_rate
+        b1, b2 = self.betas
+        eps = self.eps
+
+        m = state.get("m", mx.zeros_like(gradient))
+        v = state.get("v", mx.zeros_like(gradient))
+
+        m = b1 * m + (1 - b1) * gradient
+        v = mx.maximum(b2 * v, mx.abs(gradient))
+        state["m"] = m
+        state["v"] = v
+
+        return parameter - lr * m / (v + eps)
+
+
+class Lion(Optimizer):
+    r"""Implementation of 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.
+    We recommend a learning rate that is 3-10x smaller than AdamW and a
+    weight decay 3-10x larger than AdamW to maintain the strength
+    (lr * wd). Our Lion implementation follows the original paper. In
+    detail,
+
+    [1]: Chen, X. Symbolic Discovery of Optimization Algorithms. arXiv
+    preprint arXiv:2302.06675.
+
+    .. 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
+        w_{t + 1} &= w_t - \eta (\text{sign}(c_t) + \lambda w_t)
+
+    Args:
+        learning_rate (float): The learning rate :math:`\eta`.
+        betas (Tuple[float, float], optional): The coefficients
+          :math:`(\beta_1, \beta_2)` used for computing the gradient
+          momentum and update direction. Default: ``(0.9, 0.99)``
+        weight_decay (float, optional): The weight decay :math:`\lambda`. Default: ``0.0``
+    """
+
+    def __init__(
+        self,
+        learning_rate: float,
+        betas: List[float] = [0.9, 0.99],
+        weight_decay: float = 0.0,
+    ):
+        super().__init__()
+
+        self.learning_rate = learning_rate
+        self.betas = betas
+        self.weight_decay = weight_decay
+
+    def apply_single(
+        self, gradient: mx.array, parameter: mx.array, state: OptimizerState
+    ):
+        """Performs the Lion parameter update and stores :math:`m`
+        in the optimizer state."""
+        lr = self.learning_rate
+        b1, b2 = self.betas
+        weight_decay = self.weight_decay
+
+        m = state.get("m", gradient)
+        c = b1 * m + (1 - b1) * gradient
+        state["m"] = b2 * m + (1 - b2) * gradient
+        if weight_decay > 0:
+            parameter = (1 - lr * weight_decay) * parameter
+        return parameter - lr * mx.sign(c)

python/mlx/utils.py

@@ -1,7 +1,7 @@
 # Copyright © 2023 Apple Inc.
 
 
-def tree_map(fn, tree, *rest):
+def tree_map(fn, tree, *rest, is_leaf=None):
     """Applies ``fn`` to the leaves of the python tree ``tree`` and
     returns a new collection with the results.
 
@@ -10,6 +10,9 @@ def tree_map(fn, tree, *rest):
     ``fn``. In that respect, :meth:`tree_map` is closer to :func:`itertools.starmap`
     than to :func:`map`.
 
+    The keyword argument ``is_leaf`` decides what constitutes a leaf from
+    ``tree`` similar to :func:`tree_flatten`.
+
     .. code-block:: python
 
         import mlx.nn as nn
@@ -26,21 +29,28 @@ def tree_map(fn, tree, *rest):
         fn (Callable): The function that processes the leaves of the tree
         tree (Any): The main python tree that will be iterated upon
         rest (Tuple[Any]): Extra trees to be iterated together with tree
+        is_leaf (Optional[Callable]): An optional callable that returns True if
+            the passed object is considered a leaf or False otherwise.
 
     Returns:
         A python tree with the new values returned by ``fn``.
     """
-    if isinstance(tree, list):
+    if is_leaf is not None and is_leaf(tree):
+        return fn(tree, *rest)
+    elif isinstance(tree, list):
         return [
-            tree_map(fn, child, *(r[i] for r in rest)) for i, child in enumerate(tree)
+            tree_map(fn, child, *(r[i] for r in rest), is_leaf=is_leaf)
+            for i, child in enumerate(tree)
         ]
     elif isinstance(tree, tuple):
         return tuple(
-            tree_map(fn, child, *(r[i] for r in rest)) for i, child in enumerate(tree)
+            tree_map(fn, child, *(r[i] for r in rest), is_leaf=is_leaf)
+            for i, child in enumerate(tree)
         )
     elif isinstance(tree, dict):
         return {
-            k: tree_map(fn, child, *(r[k] for r in rest)) for k, child in tree.items()
+            k: tree_map(fn, child, *(r[k] for r in rest), is_leaf=is_leaf)
+            for k, child in tree.items()
         }
     else:
         return fn(tree, *rest)

python/src/array.cpp

@@ -436,10 +436,14 @@ void init_array(py::module_& m) {
           "__repr__",
           [](const Dtype& t) {
             std::ostringstream os;
+            os << "mlx.core.";
             os << t;
             return os.str();
           })
-      .def("__eq__", [](const Dtype& t1, const Dtype& t2) { return t1 == t2; });
+      .def("__eq__", [](const Dtype& t1, const Dtype& t2) { return t1 == t2; })
+      .def("__hash__", [](const Dtype& t) {
+        return static_cast<int64_t>(t.val);
+      });
   m.attr("bool_") = py::cast(bool_);
   m.attr("uint8") = py::cast(uint8);
   m.attr("uint16") = py::cast(uint16);
@@ -454,38 +458,54 @@ void init_array(py::module_& m) {
   m.attr("bfloat16") = py::cast(bfloat16);
   m.attr("complex64") = py::cast(complex64);
 
-  py::class_<array>(m, "array", R"pbdoc(An N-dimensional array object.)pbdoc")
-      .def(
-          py::init([](ScalarOrArray v, std::optional<Dtype> t) {
-            auto arr = to_array(v, t);
+  auto array_class = py::class_<array>(
+      m, "array", R"pbdoc(An N-dimensional array object.)pbdoc");
+
+  {
+    py::options options;
+    options.disable_function_signatures();
+
+    array_class.def(
+        py::init([](std::variant<
+                        py::bool_,
+                        py::int_,
+                        py::float_,
+                        std::complex<float>,
+                        py::list,
+                        py::tuple,
+                        py::array,
+                        py::buffer,
+                        py::object> v,
+                    std::optional<Dtype> t) {
+          if (auto pv = std::get_if<py::bool_>(&v); pv) {
+            return array(py::cast<bool>(*pv), t.value_or(bool_));
+          } else if (auto pv = std::get_if<py::int_>(&v); pv) {
+            return array(py::cast<int>(*pv), t.value_or(int32));
+          } else if (auto pv = std::get_if<py::float_>(&v); pv) {
+            return array(py::cast<float>(*pv), t.value_or(float32));
+          } else if (auto pv = std::get_if<std::complex<float>>(&v); pv) {
+            return array(static_cast<complex64_t>(*pv), t.value_or(complex64));
+          } else if (auto pv = std::get_if<py::list>(&v); pv) {
+            return array_from_list(*pv, t);
+          } else if (auto pv = std::get_if<py::tuple>(&v); pv) {
+            return array_from_list(*pv, t);
+          } else if (auto pv = std::get_if<py::array>(&v); pv) {
+            return np_array_to_mlx(*pv, t);
+          } else if (auto pv = std::get_if<py::buffer>(&v); pv) {
+            return np_array_to_mlx(*pv, t);
+          } else {
+            auto arr = to_array_with_accessor(std::get<py::object>(v));
             return astype(arr, t.value_or(arr.dtype()));
-          }),
-          "val"_a,
-          "dtype"_a = std::nullopt)
-      .def(
-          py::init([](std::variant<py::list, py::tuple> pl,
-                      std::optional<Dtype> dtype) {
-            if (auto pv = std::get_if<py::list>(&pl); pv) {
-              return array_from_list(*pv, dtype);
-            } else {
-              auto v = std::get<py::tuple>(pl);
-              return array_from_list(v, dtype);
-            }
-          }),
-          "vals"_a,
-          "dtype"_a = std::nullopt)
-      .def(
-          py::init([](py::array np_array, std::optional<Dtype> dtype) {
-            return np_array_to_mlx(np_array, dtype);
-          }),
-          "vals"_a,
-          "dtype"_a = std::nullopt)
-      .def(
-          py::init([](py::buffer np_buffer, std::optional<Dtype> dtype) {
-            return np_array_to_mlx(np_buffer, dtype);
-          }),
-          "vals"_a,
-          "dtype"_a = std::nullopt)
+          }
+        }),
+        "val"_a,
+        "dtype"_a = std::nullopt,
+        R"pbdoc(
+            __init__(self: array, val: Union[scalar, list, tuple, numpy.ndarray, array], dtype: Optional[Dtype] = None)
+          )pbdoc");
+  }
+
+  array_class
       .def_property_readonly(
           "size", &array::size, R"pbdoc(Number of elments in the array.)pbdoc")
       .def_property_readonly(
@@ -603,25 +623,51 @@ void init_array(py::module_& m) {
       .def(
           "__truediv__",
           [](const array& a, const ScalarOrArray v) {
-            return divide(a, to_array(v, float32));
+            return divide(a, to_array(v, a.dtype()));
           },
           "other"_a)
       .def(
           "__div__",
           [](const array& a, const ScalarOrArray v) {
-            return divide(a, to_array(v, float32));
+            return divide(a, to_array(v, a.dtype()));
+          },
+          "other"_a)
+      .def(
+          "__floordiv__",
+          [](const array& a, const ScalarOrArray v) {
+            auto b = to_array(v, a.dtype());
+            return floor_divide(a, b);
           },
           "other"_a)
       .def(
           "__rtruediv__",
           [](const array& a, const ScalarOrArray v) {
-            return divide(to_array(v, float32), a);
+            return divide(to_array(v, a.dtype()), a);
+          },
+          "other"_a)
+      .def(
+          "__rfloordiv__",
+          [](const array& a, const ScalarOrArray v) {
+            auto b = to_array(v, a.dtype());
+            return floor_divide(b, a);
           },
           "other"_a)
       .def(
           "__rdiv__",
           [](const array& a, const ScalarOrArray v) {
-            return divide(to_array(v, float32), a);
+            return divide(to_array(v, a.dtype()), a);
+          },
+          "other"_a)
+      .def(
+          "__mod__",
+          [](const array& a, const ScalarOrArray v) {
+            return remainder(a, to_array(v, a.dtype()));
+          },
+          "other"_a)
+      .def(
+          "__rmod__",
+          [](const array& a, const ScalarOrArray v) {
+            return remainder(to_array(v, a.dtype()), a);
           },
           "other"_a)
       .def(
@@ -680,6 +726,21 @@ void init_array(py::module_& m) {
             return power(a, to_array(v, a.dtype()));
           },
           "other"_a)
+      .def(
+          "flatten",
+          [](const array& a,
+             int start_axis,
+             int end_axis,
+             const StreamOrDevice& s) {
+            return flatten(a, start_axis, end_axis);
+          },
+          "start_axis"_a = 0,
+          "end_axis"_a = -1,
+          py::kw_only(),
+          "stream"_a = none,
+          R"pbdoc(
+            See :func:`flatten`.
+          )pbdoc")
       .def(
           "reshape",
           [](const array& a, py::args shape, StreamOrDevice s) {
@@ -814,6 +875,22 @@ void init_array(py::module_& m) {
           py::kw_only(),
           "stream"_a = none,
           "See :func:`any`.")
+      .def(
+          "moveaxis",
+          &moveaxis,
+          "source"_a,
+          "destination"_a,
+          py::kw_only(),
+          "stream"_a = none,
+          "See :func:`moveaxis`.")
+      .def(
+          "swapaxes",
+          &swapaxes,
+          "axis1"_a,
+          "axis2"_a,
+          py::kw_only(),
+          "stream"_a = none,
+          "See :func:`moveaxis`.")
       .def(
           "transpose",
           [](const array& a, py::args axes, StreamOrDevice s) {
@@ -1069,5 +1146,15 @@ void init_array(py::module_& m) {
           "reverse"_a = false,
           "inclusive"_a = true,
           "stream"_a = none,
-          "See :func:`cummin`.");
+          "See :func:`cummin`.")
+      .def(
+          "round",
+          [](const array& a, int decimals, StreamOrDevice s) {
+            return round(a, decimals, s);
+          },
+          py::pos_only(),
+          "decimals"_a = 0,
+          py::kw_only(),
+          "stream"_a = none,
+          "See :func:`round`.");
 }

python/src/indexing.cpp

@@ -41,9 +41,6 @@ void get_slice_params(
       py::getattr(in_slice, "start"), strides < 0 ? axis_size - 1 : 0);
   ends = get_slice_int(
       py::getattr(in_slice, "stop"), strides < 0 ? -axis_size - 1 : axis_size);
-
-  // starts = (starts < 0) ? starts + axis_size : starts;
-  // ends = (ends < 0) ? ends + axis_size : ends;
 }
 
 array get_int_index(py::object idx, int axis_size) {
@@ -123,6 +120,11 @@ array mlx_gather_nd(
     if (py::isinstance<py::slice>(idx)) {
       int start, end, stride;
       get_slice_params(start, end, stride, idx, src.shape(i));
+
+      // Handle negative indices
+      start = (start < 0) ? start + src.shape(i) : start;
+      end = (end < 0) ? end + src.shape(i) : end;
+
       gather_indices.push_back(arange(start, end, stride, uint32));
       num_slices++;
       is_slice[i] = true;
@@ -568,7 +570,13 @@ array mlx_set_item_nd(
     auto& pyidx = indices[i];
     if (py::isinstance<py::slice>(pyidx)) {
       int start, end, stride;
-      get_slice_params(start, end, stride, pyidx, src.shape(ax++));
+      auto axis_size = src.shape(ax++);
+      get_slice_params(start, end, stride, pyidx, axis_size);
+
+      // Handle negative indices
+      start = (start < 0) ? start + axis_size : start;
+      end = (end < 0) ? end + axis_size : end;
+
       auto idx = arange(start, end, stride, uint32);
       std::vector<int> idx_shape(max_dim + num_slices, 1);
       auto loc = slice_num + (arrays_first ? max_dim : 0);

python/src/load.cpp

@@ -11,8 +11,6 @@
 #include <unordered_map>
 #include <vector>
 
-#include <iostream>
-
 #include "mlx/load.h"
 #include "mlx/ops.h"
 #include "mlx/utils.h"
@@ -28,7 +26,7 @@ using namespace mlx::core;
 ///////////////////////////////////////////////////////////////////////////////
 
 bool is_istream_object(const py::object& file) {
-  return py::hasattr(file, "read") && py::hasattr(file, "seek") &&
+  return py::hasattr(file, "readinto") && py::hasattr(file, "seek") &&
       py::hasattr(file, "tell") && py::hasattr(file, "closed");
 }
 
@@ -99,26 +97,54 @@ class PyFileReader : public io::Reader {
         seek_func_(file.attr("seek")),
         tell_func_(file.attr("tell")) {}
 
+  ~PyFileReader() {
+    py::gil_scoped_acquire gil;
+
+    pyistream_.release().dec_ref();
+    readinto_func_.release().dec_ref();
+    seek_func_.release().dec_ref();
+    tell_func_.release().dec_ref();
+  }
+
   bool is_open() const override {
-    return !pyistream_.attr("closed").cast<bool>();
+    bool out;
+    {
+      py::gil_scoped_acquire gil;
+      out = !pyistream_.attr("closed").cast<bool>();
+    }
+    return out;
   }
 
   bool good() const override {
-    return !pyistream_.is_none();
+    bool out;
+    {
+      py::gil_scoped_acquire gil;
+      out = !pyistream_.is_none();
+    }
+    return out;
   }
 
   size_t tell() const override {
-    return tell_func_().cast<size_t>();
+    size_t out;
+    {
+      py::gil_scoped_acquire gil;
+      out = tell_func_().cast<size_t>();
+    }
+    return out;
   }
 
   void seek(int64_t off, std::ios_base::seekdir way = std::ios_base::beg)
       override {
+    py::gil_scoped_acquire gil;
     seek_func_(off, (int)way);
   }
 
   void read(char* data, size_t n) override {
+    py::gil_scoped_acquire gil;
+
     py::object bytes_read =
         readinto_func_(py::memoryview::from_buffer(data, {n}, {sizeof(char)}));
+
     if (bytes_read.is_none() || py::cast<size_t>(bytes_read) < n) {
       throw std::runtime_error("[load] Failed to read from python stream");
     }
@@ -163,6 +189,7 @@ DictOrArray mlx_load_helper(py::object file, StreamOrDevice s) {
 
     // If we don't own the stream and it was passed to us, eval immediately
     for (auto& [key, arr] : array_dict) {
+      py::gil_scoped_release gil;
       arr.eval();
     }
 
@@ -172,7 +199,10 @@ DictOrArray mlx_load_helper(py::object file, StreamOrDevice s) {
   } else if (is_istream_object(file)) {
     // If we don't own the stream and it was passed to us, eval immediately
     auto arr = load(std::make_shared<PyFileReader>(file), s);
-    arr.eval();
+    {
+      py::gil_scoped_release gil;
+      arr.eval();
+    }
     return {arr};
   }
 
@@ -192,26 +222,54 @@ class PyFileWriter : public io::Writer {
         seek_func_(file.attr("seek")),
         tell_func_(file.attr("tell")) {}
 
+  ~PyFileWriter() {
+    py::gil_scoped_acquire gil;
+
+    pyostream_.release().dec_ref();
+    write_func_.release().dec_ref();
+    seek_func_.release().dec_ref();
+    tell_func_.release().dec_ref();
+  }
+
   bool is_open() const override {
-    return !pyostream_.attr("closed").cast<bool>();
+    bool out;
+    {
+      py::gil_scoped_acquire gil;
+      out = !pyostream_.attr("closed").cast<bool>();
+    }
+    return out;
   }
 
   bool good() const override {
-    return !pyostream_.is_none();
+    bool out;
+    {
+      py::gil_scoped_acquire gil;
+      out = !pyostream_.is_none();
+    }
+    return out;
   }
 
   size_t tell() const override {
-    return tell_func_().cast<size_t>();
+    size_t out;
+    {
+      py::gil_scoped_acquire gil;
+      out = tell_func_().cast<size_t>();
+    }
+    return out;
   }
 
   void seek(int64_t off, std::ios_base::seekdir way = std::ios_base::beg)
       override {
+    py::gil_scoped_acquire gil;
     seek_func_(off, (int)way);
   }
 
   void write(const char* data, size_t n) override {
+    py::gil_scoped_acquire gil;
+
     py::object bytes_written =
         write_func_(py::memoryview::from_buffer(data, {n}, {sizeof(char)}));
+
     if (bytes_written.is_none() || py::cast<size_t>(bytes_written) < n) {
       throw std::runtime_error("[load] Failed to write to python stream");
     }
@@ -228,12 +286,21 @@ class PyFileWriter : public io::Writer {
   py::object tell_func_;
 };
 
-void mlx_save_helper(py::object file, array a, bool retain_graph) {
+void mlx_save_helper(
+    py::object file,
+    array a,
+    std::optional<bool> retain_graph_) {
+  bool retain_graph = retain_graph_.value_or(a.is_tracer());
   if (py::isinstance<py::str>(file)) {
     save(py::cast<std::string>(file), a, retain_graph);
     return;
   } else if (is_ostream_object(file)) {
-    save(std::make_shared<PyFileWriter>(file), a, retain_graph);
+    auto writer = std::make_shared<PyFileWriter>(file);
+    {
+      py::gil_scoped_release gil;
+      save(writer, a, retain_graph);
+    }
+
     return;
   }
 
@@ -285,7 +352,11 @@ void mlx_savez_helper(
   for (auto [k, a] : arrays_dict) {
     std::string fname = k + ".npy";
     auto py_ostream = zipfile_object.open(fname, 'w');
-    save(std::make_shared<PyFileWriter>(py_ostream), a);
+    auto writer = std::make_shared<PyFileWriter>(py_ostream);
+    {
+      py::gil_scoped_release gil;
+      save(writer, a, /*retain_graph=*/a.is_tracer());
+    }
   }
 
   return;

python/src/load.h

@@ -13,9 +13,12 @@ using namespace mlx::core;
 using DictOrArray = std::variant<array, std::unordered_map<std::string, array>>;
 
 DictOrArray mlx_load_helper(py::object file, StreamOrDevice s);
-void mlx_save_helper(py::object file, array a, bool retain_graph = true);
+void mlx_save_helper(
+    py::object file,
+    array a,
+    std::optional<bool> retain_graph = std::nullopt);
 void mlx_savez_helper(
     py::object file,
     py::args args,
     const py::kwargs& kwargs,
-    bool compressed = false);
+    bool compressed = false);

python/src/mlx.cpp

@@ -17,7 +17,7 @@ void init_random(py::module_&);
 void init_fft(py::module_&);
 
 PYBIND11_MODULE(core, m) {
-  m.doc() = "mlx: A framework for machine learning on Apple Silicon.";
+  m.doc() = "mlx: A framework for machine learning on Apple silicon.";
 
   auto reprlib_fix = py::module_::import("mlx._reprlib_fix");
 

python/src/utils.h

@@ -15,8 +15,8 @@ namespace py = pybind11;
 using namespace mlx::core;
 
 using IntOrVec = std::variant<std::monostate, int, std::vector<int>>;
-using ScalarOrArray =
-    std::variant<py::bool_, py::int_, py::float_, std::complex<float>, array>;
+using ScalarOrArray = std::
+    variant<py::bool_, py::int_, py::float_, std::complex<float>, py::object>;
 static constexpr std::monostate none{};
 
 inline std::vector<int> get_reduce_axes(const IntOrVec& v, int dims) {
@@ -32,6 +32,14 @@ inline std::vector<int> get_reduce_axes(const IntOrVec& v, int dims) {
   return axes;
 }
 
+inline array to_array_with_accessor(py::object obj) {
+  if (py::hasattr(obj, "__mlx_array__")) {
+    return obj.attr("__mlx_array__")().cast<array>();
+  } else {
+    return obj.cast<array>();
+  }
+}
+
 inline array to_array(
     const ScalarOrArray& v,
     std::optional<Dtype> dtype = std::nullopt) {
@@ -48,7 +56,7 @@ inline array to_array(
   } else if (auto pv = std::get_if<std::complex<float>>(&v); pv) {
     return array(static_cast<complex64_t>(*pv), complex64);
   } else {
-    return std::get<array>(v);
+    return to_array_with_accessor(std::get<py::object>(v));
   }
 }
 
@@ -60,13 +68,16 @@ inline std::pair<array, array> to_arrays(
   // - If a is an array but b is not, treat b as a weak python type
   // - If b is an array but a is not, treat a as a weak python type
   // - If neither is an array convert to arrays but leave their types alone
-  if (auto pa = std::get_if<array>(&a); pa) {
-    if (auto pb = std::get_if<array>(&b); pb) {
-      return {*pa, *pb};
+  if (auto pa = std::get_if<py::object>(&a); pa) {
+    auto arr_a = to_array_with_accessor(*pa);
+    if (auto pb = std::get_if<py::object>(&b); pb) {
+      auto arr_b = to_array_with_accessor(*pb);
+      return {arr_a, arr_b};
     }
-    return {*pa, to_array(b, pa->dtype())};
-  } else if (auto pb = std::get_if<array>(&b); pb) {
-    return {to_array(a, pb->dtype()), *pb};
+    return {arr_a, to_array(b, arr_a.dtype())};
+  } else if (auto pb = std::get_if<py::object>(&b); pb) {
+    auto arr_b = to_array_with_accessor(*pb);
+    return {to_array(a, arr_b.dtype()), arr_b};
   } else {
     return {to_array(a), to_array(b)};
   }

python/tests/mlx_tests.py

@@ -4,6 +4,7 @@ import os
 import unittest
 
 import mlx.core as mx
+import numpy as np
 
 
 class MLXTestCase(unittest.TestCase):
@@ -16,3 +17,19 @@ class MLXTestCase(unittest.TestCase):
 
     def tearDown(self):
         mx.set_default_device(self.default)
+
+    def assertEqualArray(
+        self,
+        mx_res: mx.array,
+        expected: mx.array,
+        atol=1e-2,
+        rtol=1e-2,
+        **kwargs,
+    ):
+        assert tuple(mx_res.shape) == tuple(
+            expected.shape
+        ), f"shape mismatch expected={expected.shape} got={mx_res.shape}"
+        assert (
+            mx_res.dtype == expected.dtype
+        ), f"dtype mismatch expected={expected.dtype} got={mx_res.dtype}"
+        np.testing.assert_allclose(mx_res, expected, rtol=rtol, atol=atol)

python/tests/test_array.py

@@ -5,9 +5,8 @@ import unittest
 from itertools import permutations
 
 import mlx.core as mx
-import numpy as np
-
 import mlx_tests
+import numpy as np
 
 
 class TestVersion(mlx_tests.MLXTestCase):
@@ -35,19 +34,19 @@ class TestDtypes(mlx_tests.MLXTestCase):
         self.assertEqual(mx.bfloat16.size, 2)
         self.assertEqual(mx.complex64.size, 8)
 
-        self.assertEqual(str(mx.bool_), "bool")
-        self.assertEqual(str(mx.uint8), "uint8")
-        self.assertEqual(str(mx.uint16), "uint16")
-        self.assertEqual(str(mx.uint32), "uint32")
-        self.assertEqual(str(mx.uint64), "uint64")
-        self.assertEqual(str(mx.int8), "int8")
-        self.assertEqual(str(mx.int16), "int16")
-        self.assertEqual(str(mx.int32), "int32")
-        self.assertEqual(str(mx.int64), "int64")
-        self.assertEqual(str(mx.float16), "float16")
-        self.assertEqual(str(mx.float32), "float32")
-        self.assertEqual(str(mx.bfloat16), "bfloat16")
-        self.assertEqual(str(mx.complex64), "complex64")
+        self.assertEqual(str(mx.bool_), "mlx.core.bool")
+        self.assertEqual(str(mx.uint8), "mlx.core.uint8")
+        self.assertEqual(str(mx.uint16), "mlx.core.uint16")
+        self.assertEqual(str(mx.uint32), "mlx.core.uint32")
+        self.assertEqual(str(mx.uint64), "mlx.core.uint64")
+        self.assertEqual(str(mx.int8), "mlx.core.int8")
+        self.assertEqual(str(mx.int16), "mlx.core.int16")
+        self.assertEqual(str(mx.int32), "mlx.core.int32")
+        self.assertEqual(str(mx.int64), "mlx.core.int64")
+        self.assertEqual(str(mx.float16), "mlx.core.float16")
+        self.assertEqual(str(mx.float32), "mlx.core.float32")
+        self.assertEqual(str(mx.bfloat16), "mlx.core.bfloat16")
+        self.assertEqual(str(mx.complex64), "mlx.core.complex64")
 
     def test_scalar_conversion(self):
         dtypes = [
@@ -728,6 +727,11 @@ class TestArray(mlx_tests.MLXTestCase):
             np.array_equal(a_np[idx_np, idx_np], np.array(a_mlx[idx_mlx, idx_mlx]))
         )
 
+        # Slicing with negative indices and integer
+        a_np = np.arange(10).reshape(5, 2)
+        a_mlx = mx.array(a_np)
+        self.assertTrue(np.array_equal(a_np[2:-1, 0], np.array(a_mlx[2:-1, 0])))
+
     def test_setitem(self):
         a = mx.array(0)
         a[None] = 1
@@ -904,6 +908,11 @@ class TestArray(mlx_tests.MLXTestCase):
             np.array([0, 1]),
         )
 
+        # Check slice assign with negative indices works
+        a = mx.zeros((5, 5), mx.int32)
+        a[2:-2, 2:-2] = 4
+        self.assertEqual(a[2, 2].item(), 4)
+
     def test_slice_negative_step(self):
         a_np = np.arange(20)
         a_mx = mx.array(a_np)

python/tests/test_autograd.py

@@ -3,7 +3,6 @@
 import unittest
 
 import mlx.core as mx
-
 import mlx_tests
 
 

python/tests/test_bf16.py

@@ -1,13 +1,12 @@
 # Copyright © 2023 Apple Inc.
 
+import math
 import unittest
 from itertools import permutations
 
-import math
 import mlx.core as mx
-import numpy as np
-
 import mlx_tests
+import numpy as np
 
 try:
     import torch