bump (#1931)

* do affine quant in fp32

* static cast

CMakeLists.txt

@@ -1,6 +1,23 @@
 cmake_minimum_required(VERSION 3.25)
 
-project(mlx LANGUAGES C CXX)
+if(NOT MLX_VERSION)
+  file(STRINGS "mlx/version.h" _mlx_h_version REGEX "^#define MLX_VERSION_.*$")
+  string(REGEX MATCH "#define MLX_VERSION_MAJOR ([0-9]+)" _ "${_mlx_h_version}")
+  set(_major ${CMAKE_MATCH_1})
+  string(REGEX MATCH "#define MLX_VERSION_MINOR ([0-9]+)" _ "${_mlx_h_version}")
+  set(_minor ${CMAKE_MATCH_1})
+  string(REGEX MATCH "#define MLX_VERSION_PATCH ([0-9]+)" _ "${_mlx_h_version}")
+  set(_patch ${CMAKE_MATCH_1})
+  set(MLX_PROJECT_VERSION "${_major}.${_minor}.${_patch}")
+else()
+  string(REGEX REPLACE "^([0-9]+\.[0-9]+\.[0-9]+).*" "\\1" MLX_PROJECT_VERSION
+                       ${MLX_VERSION})
+endif()
+
+project(
+  mlx
+  LANGUAGES C CXX
+  VERSION ${MLX_PROJECT_VERSION})
 
 # ----------------------------- Setup -----------------------------
 set(CMAKE_MODULE_PATH "${PROJECT_SOURCE_DIR}/cmake")
@@ -24,13 +41,9 @@ option(MLX_BUILD_BLAS_FROM_SOURCE "Build OpenBLAS from source code" OFF)
 option(MLX_METAL_JIT "Use JIT compilation for Metal kernels" OFF)
 option(BUILD_SHARED_LIBS "Build mlx as a shared library" OFF)
 
-if(NOT MLX_VERSION)
-  set(MLX_VERSION 0.23.0)
-endif()
 add_compile_definitions("MLX_VERSION=${MLX_VERSION}")
 
 # --------------------- Processor tests -------------------------
-
 message(
   STATUS
     "Building MLX for ${CMAKE_SYSTEM_PROCESSOR} processor on ${CMAKE_SYSTEM_NAME}"
@@ -64,6 +77,7 @@ include(FetchContent)
 cmake_policy(SET CMP0135 NEW)
 
 add_library(mlx)
+set_target_properties(mlx PROPERTIES COMPILE_WARNING_AS_ERROR ON)
 
 if(MLX_BUILD_METAL)
   set(METAL_LIB "-framework Metal")

benchmarks/python/layer_norm_bench.py

@@ -10,7 +10,12 @@ def layer_norm(x, w, b, eps):
     x = x.astype(mx.float32)
     mu = mx.mean(x, -1, keepdims=True)
     v = mx.var(x, -1, keepdims=True)
-    return (x - mu) * mx.rsqrt(v + eps) * w + b
+    y = (x - mu) * mx.rsqrt(v + eps)
+    if w is not None:
+        y = y * w
+    if b is not None:
+        y = y + b
+    return y
 
 
 def time_layer_norm():
@@ -36,6 +41,28 @@ def time_layer_norm():
     time_fn(layer_norm_loop, mx.compile(g1), x, w, b)
     time_fn(layer_norm_loop, mx.compile(g2), x, w, b)
 
+    f1 = lambda x, y: (layer_norm(x, None, None, 1e-5) * y).sum()
+    f2 = lambda x, y: (mx.fast.layer_norm(x, None, None, 1e-5) * y).sum()
+    g1 = mx.grad(f1, argnums=(0,))
+    g2 = mx.grad(f2, argnums=(0,))
+
+    x = mx.random.uniform(shape=(8, 1024, 4096)).astype(mx.float16)
+    w = mx.random.uniform(shape=(4096,)).astype(mx.float16)
+    b = mx.random.uniform(shape=(4096,)).astype(mx.float16)
+    y = mx.random.uniform(shape=(8, 1024, 4096)).astype(mx.float16)
+    mx.eval(x, w, b, y)
+
+    def layer_norm_loop(g, x):
+        gx = x
+        for _ in range(32):
+            gx = g(gx, y)
+        return gx
+
+    time_fn(layer_norm_loop, g1, x)
+    time_fn(layer_norm_loop, g2, x)
+    time_fn(layer_norm_loop, mx.compile(g1), x)
+    time_fn(layer_norm_loop, mx.compile(g2), x)
+
 
 if __name__ == "__main__":
     time_layer_norm()

benchmarks/python/rms_norm_bench.py

@@ -9,7 +9,10 @@ def rms_norm(x, w, eps):
     ot = x.dtype
     x = x.astype(mx.float32)
     n = mx.rsqrt(x.square().mean(-1, keepdims=True) + eps)
-    return (x * n).astype(ot) * w
+    y = (x * n).astype(ot)
+    if w is not None:
+        y = y * w
+    return y
 
 
 def time_rms_norm():
@@ -34,6 +37,27 @@ def time_rms_norm():
     time_fn(rms_norm_loop, mx.compile(g1), x, w)
     time_fn(rms_norm_loop, mx.compile(g2), x, w)
 
+    f1 = lambda x, y: (rms_norm(x, None, 1e-5) * y).sum()
+    f2 = lambda x, y: (mx.fast.rms_norm(x, None, 1e-5) * y).sum()
+    g1 = mx.grad(f1, argnums=(0,))
+    g2 = mx.grad(f2, argnums=(0,))
+
+    x = mx.random.uniform(shape=(8, 1024, 4096)).astype(mx.float16)
+    w = mx.random.uniform(shape=(4096,)).astype(mx.float16)
+    y = mx.random.uniform(shape=(8, 1024, 4096)).astype(mx.float16)
+    mx.eval(x, w, y)
+
+    def rms_norm_loop(g, x):
+        gx = x
+        for _ in range(32):
+            gx = g(gx, y)
+        return gx
+
+    time_fn(rms_norm_loop, g1, x)
+    time_fn(rms_norm_loop, g2, x)
+    time_fn(rms_norm_loop, mx.compile(g1), x)
+    time_fn(rms_norm_loop, mx.compile(g2), x)
+
 
 if __name__ == "__main__":
     time_rms_norm()

cmake/extension.cmake

@@ -1,5 +1,7 @@
 include(CMakeParseArguments)
 
+# clang format off
+#
 # ##############################################################################
 # Build metal library
 #
@@ -11,6 +13,8 @@ include(CMakeParseArguments)
 # of source files INCLUDE_DIRS: List of include dirs DEPS: List of dependency
 # files (like headers)
 #
+# clang format on
+
 macro(mlx_build_metallib)
   # Parse args
   set(oneValueArgs TARGET TITLE OUTPUT_DIRECTORY)
@@ -21,7 +25,7 @@ macro(mlx_build_metallib)
   set(MTLLIB_BUILD_TARGET "${MTLLIB_OUTPUT_DIRECTORY}/${MTLLIB_TITLE}.metallib")
 
   # Collect compile options
-  set(MTLLIB_COMPILE_OPTIONS -Wall -Wextra -fno-fast-math)
+  set(MTLLIB_COMPILE_OPTIONS -Wall -Wextra -fno-fast-math -Wno-c++17-extensions)
 
   # Prepare metallib build command
   add_custom_command(

docs/src/python/nn.rst

@@ -174,6 +174,7 @@ In detail:
 
    value_and_grad
    quantize
+   average_gradients
 
 .. toctree::
 

docs/src/usage/distributed.rst

@@ -5,21 +5,27 @@ Distributed Communication
 
 .. currentmodule:: mlx.core.distributed
 
-MLX utilizes `MPI <https://en.wikipedia.org/wiki/Message_Passing_Interface>`_ to
-provide distributed communication operations that allow the computational cost
-of training or inference to be shared across many physical machines. You can
-see a list of the supported operations in the :ref:`API docs<distributed>`.
+MLX supports distributed communication operations that allow the computational cost
+of training or inference to be shared across many physical machines. At the
+moment we support two different communication backends:
+
+* `MPI <https://en.wikipedia.org/wiki/Message_Passing_Interface>`_ a
+  full-featured and mature distributed communications library
+* A **ring** backend of our own that uses native TCP sockets and should be
+  faster for thunderbolt connections.
+
+The list of all currently supported operations and their documentation can be
+seen in the :ref:`API docs<distributed>`.
 
 .. note::
-   A lot of operations may not be supported or not as fast as they should be.
+   Some operations may not be supported or not as fast as they should be.
    We are adding more and tuning the ones we have as we are figuring out the
    best way to do distributed computing on Macs using MLX.
 
 Getting Started
 ---------------
 
-MLX already comes with the ability to "talk" to MPI if it is installed on the
-machine. The minimal distributed program in MLX is as simple as:
+A distributed program in MLX is as simple as:
 
 .. code:: python
 
@@ -30,74 +36,79 @@ machine. The minimal distributed program in MLX is as simple as:
     print(world.rank(), x)
 
 The program above sums the array ``mx.ones(10)`` across all
-distributed processes. If simply run with ``python``, however, only one
-process is launched and no distributed communication takes place.
+distributed processes. However, when this script is run with ``python`` only
+one process is launched and no distributed communication takes place. Namely,
+all operations in ``mx.distributed`` are noops when the distributed group has a
+size of one. This property allows us to avoid code that checks if we are in a
+distributed setting similar to the one below:
 
-To launch the program in distributed mode we need to use ``mpirun`` or
-``mpiexec`` depending on the MPI installation. The simplest possible way is the
-following:
+.. code:: python
+
+    import mlx.core as mx
+
+    x = ...
+    world = mx.distributed.init()
+    # No need for the check we can simply do x = mx.distributed.all_sum(x)
+    if world.size() > 1:
+        x = mx.distributed.all_sum(x)
+
+Running Distributed Programs
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+MLX provides ``mlx.launch`` a helper script to launch distributed programs.
+Continuing with our initial example we can run it on localhost with 4 processes using
 
 .. code:: shell
 
-    $ mpirun -np 2 python test.py
-    1 array([2, 2, 2, ..., 2, 2, 2], dtype=float32)
-    0 array([2, 2, 2, ..., 2, 2, 2], dtype=float32)
+    $ mlx.launch -n 4 my_script.py
+    3 array([4, 4, 4, ..., 4, 4, 4], dtype=float32)
+    2 array([4, 4, 4, ..., 4, 4, 4], dtype=float32)
+    1 array([4, 4, 4, ..., 4, 4, 4], dtype=float32)
+    0 array([4, 4, 4, ..., 4, 4, 4], dtype=float32)
 
-The above launches two processes on the same (local) machine and we can see
-both standard output streams. The processes send the array of 1s to each other
-and compute the sum which is printed. Launching with ``mpirun -np 4 ...`` would
-print 4 etc.
-
-Installing MPI
----------------
-
-MPI can be installed with Homebrew, using the Anaconda package manager or
-compiled from source. Most of our testing is done using ``openmpi`` installed
-with the Anaconda package manager as follows:
+We can also run it on some remote hosts by providing their IPs (provided that
+the script exists on all hosts and they are reachable by ssh)
 
 .. code:: shell
 
-    $ conda install conda-forge::openmpi
+    $ mlx.launch --hosts ip1,ip2,ip3,ip4 my_script.py
+    3 array([4, 4, 4, ..., 4, 4, 4], dtype=float32)
+    2 array([4, 4, 4, ..., 4, 4, 4], dtype=float32)
+    1 array([4, 4, 4, ..., 4, 4, 4], dtype=float32)
+    0 array([4, 4, 4, ..., 4, 4, 4], dtype=float32)
 
-Installing with Homebrew may require specifying the location of ``libmpi.dyld``
-so that MLX can find it and load it at runtime. This can simply be achieved by
-passing the ``DYLD_LIBRARY_PATH`` environment variable to ``mpirun``.
+Consult the dedicated :doc:`usage guide<launching_distributed>` for more
+information on using ``mlx.launch``.
 
-.. code:: shell
+Selecting Backend
+^^^^^^^^^^^^^^^^^
 
-    $ mpirun -np 2 -x DYLD_LIBRARY_PATH=/opt/homebrew/lib/ python test.py
-
-Setting up Remote Hosts
------------------------
-
-MPI can automatically connect to remote hosts and set up the communication over
-the network if the remote hosts can be accessed via ssh. A good checklist to
-debug connectivity issues is the following:
-
-* ``ssh hostname`` works from all machines to all machines without asking for
-  password or host confirmation
-* ``mpirun`` is accessible on all machines. You can call ``mpirun`` using its
-  full path to force all machines to use a specific path.
-* Ensure that the ``hostname`` used by MPI is the one that you have configured
-  in the ``.ssh/config`` files on all machines.
+You can select the backend you want to use when calling :func:`init` by passing
+one of ``{'any', 'ring', 'mpi'}``. When passing ``any``, MLX will try to
+initialize the ``ring`` backend and if it fails the ``mpi`` backend. If they
+both fail then a singleton group is created.
 
 .. note::
-  For an example hostname ``foo.bar.com`` MPI can use only ``foo`` as
-  the hostname passed to ssh if the current hostname matches ``*.bar.com``.
+   After a distributed backend is successfully initialized :func:`init` will
+   return **the same backend** if called without arguments or with backend set to
+   ``any``.
 
-An easy way to pass the host names to MPI is using a host file. A host file
-looks like the following, where ``host1`` and ``host2`` should be the fully
-qualified domain names or IPs for these hosts.
+The following examples aim to clarify the backend initialization logic in MLX:
 
-.. code::
+.. code:: python
 
-    host1 slots=1
-    host2 slots=1
+    # Case 1: Initialize MPI regardless if it was possible to initialize the ring backend
+    world = mx.distributed.init(backend="mpi")
+    world2 = mx.distributed.init()  # subsequent calls return the MPI backend!
 
-When using MLX, it is very likely that you want to use 1 slot per host, ie one
-process per host.  The hostfile also needs to contain the current
-host if you want to run on the local host. Passing the host file to
-``mpirun`` is simply done using the ``--hostfile`` command line argument.
+    # Case 2: Initialize any backend
+    world = mx.distributed.init(backend="any")  # equivalent to no arguments
+    world2 = mx.distributed.init()  # same as above
+
+    # Case 3: Initialize both backends at the same time
+    world_mpi = mx.distributed.init(backend="mpi")
+    world_ring = mx.distributed.init(backend="ring")
+    world_any = mx.distributed.init()  # same as MPI because it was initialized first!
 
 Training Example
 ----------------
@@ -155,13 +166,179 @@ everything else remaining the same.
         optimizer.update(model, grads)
         return loss
 
-Tuning All Reduce
------------------
+Utilizing ``nn.average_gradients``
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
-We are working on improving the performance of all reduce on MLX but for now
-the two main things one can do to extract the most out of distributed training with MLX are:
+Although the code example above works correctly; it performs one communication
+per gradient. It is significantly more efficient to aggregate several gradients
+together and perform fewer communication steps.
 
-1. Perform a few large reductions instead of many small ones to improve
-   bandwidth and latency
-2. Pass ``--mca btl_tcp_links 4`` to ``mpirun`` to configure it to use 4 tcp
-   connections between each host to improve bandwidth
+This is the purpose of :func:`mlx.nn.average_gradients`. The final code looks
+almost identical to the example above:
+
+.. code:: python
+
+    model = ...
+    optimizer = ...
+    dataset = ...
+
+    def step(model, x, y):
+        loss, grads = loss_grad_fn(model, x, y)
+        grads = mlx.nn.average_gradients(grads) # <---- This line was added
+        optimizer.update(model, grads)
+        return loss
+
+    for x, y in dataset:
+        loss = step(model, x, y)
+        mx.eval(loss, model.parameters())
+
+
+Getting Started with MPI
+------------------------
+
+MLX already comes with the ability to "talk" to MPI if it is installed on the
+machine. Launching distributed MLX programs that use MPI can be done with
+``mpirun`` as expected. However, in the following examples we will be using
+``mlx.launch --backend mpi`` which takes care of some nuisances such as setting
+absolute paths for the ``mpirun`` executable and the ``libmpi.dyld`` shared
+library.
+
+The simplest possible usage is the following which, assuming the minimal
+example in the beginning of this page, should result in:
+
+.. code:: shell
+
+    $ mlx.launch --backend mpi -n 2 test.py
+    1 array([2, 2, 2, ..., 2, 2, 2], dtype=float32)
+    0 array([2, 2, 2, ..., 2, 2, 2], dtype=float32)
+
+The above launches two processes on the same (local) machine and we can see
+both standard output streams. The processes send the array of 1s to each other
+and compute the sum which is printed. Launching with ``mlx.launch -n 4 ...`` would
+print 4 etc.
+
+Installing MPI
+^^^^^^^^^^^^^^
+
+MPI can be installed with Homebrew, using the Anaconda package manager or
+compiled from source. Most of our testing is done using ``openmpi`` installed
+with the Anaconda package manager as follows:
+
+.. code:: shell
+
+    $ conda install conda-forge::openmpi
+
+Installing with Homebrew may require specifying the location of ``libmpi.dyld``
+so that MLX can find it and load it at runtime. This can simply be achieved by
+passing the ``DYLD_LIBRARY_PATH`` environment variable to ``mpirun`` and it is
+done automatically by ``mlx.launch``.
+
+.. code:: shell
+
+    $ mpirun -np 2 -x DYLD_LIBRARY_PATH=/opt/homebrew/lib/ python test.py
+    $ # or simply
+    $ mlx.launch -n 2 test.py
+
+Setting up Remote Hosts
+^^^^^^^^^^^^^^^^^^^^^^^
+
+MPI can automatically connect to remote hosts and set up the communication over
+the network if the remote hosts can be accessed via ssh. A good checklist to
+debug connectivity issues is the following:
+
+* ``ssh hostname`` works from all machines to all machines without asking for
+  password or host confirmation
+* ``mpirun`` is accessible on all machines.
+* Ensure that the ``hostname`` used by MPI is the one that you have configured
+  in the ``.ssh/config`` files on all machines.
+
+Tuning MPI All Reduce
+^^^^^^^^^^^^^^^^^^^^^
+
+.. note::
+
+    For faster all reduce consider using the ring backend either with Thunderbolt
+    connections or over Ethernet.
+
+Configure MPI to use N tcp connections between each host to improve bandwidth
+by passing ``--mca btl_tcp_links N``.
+
+Force MPI to use the most performant network interface by setting ``--mca
+btl_tcp_if_include <iface>`` where ``<iface>`` should be the interface you want
+to use.
+
+Getting Started with Ring
+-------------------------
+
+The ring backend does not depend on any third party library so it is always
+available. It uses TCP sockets so the nodes need to be reachable via a network.
+As the name suggests the nodes are connected in a ring which means that rank 1
+can only communicate with rank 0 and rank 2, rank 2 only with rank 1 and rank 3
+and so on and so forth. As a result :func:`send` and :func:`recv` with
+arbitrary sender and receiver is not supported in the ring backend.
+
+Defining a Ring
+^^^^^^^^^^^^^^^
+
+The easiest way to define and use a ring is via a JSON hostfile and the
+``mlx.launch`` :doc:`helper script <launching_distributed>`. For each node one
+defines a hostname to ssh into to run commands on this node and one or more IPs
+that this node will listen to for connections.
+
+For example the hostfile below defines a 4 node ring. ``hostname1`` will be
+rank 0, ``hostname2`` rank 1 etc.
+
+.. code:: json
+
+    [
+        {"ssh": "hostname1", "ips": ["123.123.123.1"]},
+        {"ssh": "hostname2", "ips": ["123.123.123.2"]},
+        {"ssh": "hostname3", "ips": ["123.123.123.3"]},
+        {"ssh": "hostname4", "ips": ["123.123.123.4"]}
+    ]
+
+Running ``mlx.launch --hostfile ring-4.json my_script.py`` will ssh into each
+node, run the script which will listen for connections in each of the provided
+IPs. Specifically, ``hostname1`` will connect to ``123.123.123.2`` and accept a
+connection from ``123.123.123.4`` and so on and so forth.
+
+Thunderbolt Ring
+^^^^^^^^^^^^^^^^
+
+Although the ring backend can have benefits over MPI even for Ethernet, its
+main purpose is to use Thunderbolt rings for higher bandwidth communication.
+Setting up such thunderbolt rings can be done manually, but is a relatively
+tedious process. To simplify this, we provide the utility ``mlx.distributed_config``.
+
+To use ``mlx.distributed_config`` your computers need to be accessible by ssh via
+Ethernet or Wi-Fi. Subsequently, connect them via thunderbolt cables and then call the
+utility as follows:
+
+.. code:: shell
+
+   mlx.distributed_config --verbose --hosts host1,host2,host3,host4
+
+By default the script will attempt to discover the thunderbolt ring and provide
+you with the commands to configure each node as well as the ``hostfile.json``
+to use with ``mlx.launch``. If password-less ``sudo`` is available on the nodes
+then ``--auto-setup`` can be used to configure them automatically.
+
+To validate your connection without configuring anything
+``mlx.distributed_config`` can also plot the ring using DOT format.
+
+.. code:: shell
+
+   mlx.distributed_config --verbose --hosts host1,host2,host3,host4 --dot >ring.dot
+   dot -Tpng ring.dot >ring.png
+   open ring.png
+
+If you want to go through the process manually, the steps are as follows:
+
+* Disable the thunderbolt bridge interface
+* For the cable connecting rank ``i`` to rank ``i + 1`` find the interfaces
+  corresponding to that cable in nodes ``i`` and ``i + 1``.
+* Set up a unique subnetwork connecting the two nodes for the corresponding
+  interfaces. For instance if the cable corresponds to ``en2`` on node ``i``
+  and ``en2`` also on node ``i + 1`` then we may assign IPs ``192.168.0.1`` and
+  ``192.168.0.2`` respectively to the two nodes. For more details you can see
+  the commands prepared by the utility script.

docs/src/usage/launching_distributed.rst

@@ -0,0 +1,105 @@
+:orphan:
+
+.. _usage_launch_distributed:
+
+Launching Distributed Programs
+==============================
+
+.. currentmodule:: mlx.core.distributed
+
+Installing the MLX python package provides a helper script ``mlx.launch`` that
+can be used to run python scripts distributed on several nodes. It allows
+launching using either the MPI backend or the ring backend. See the
+:doc:`distributed docs <distributed>` for the different backends.
+
+Usage
+-----
+
+The minimal usage example of ``mlx.launch`` is simply
+
+.. code:: shell
+
+    mlx.launch --hosts ip1,ip2 my_script.py
+
+or for testing on localhost
+
+.. code:: shell
+
+    mlx.launch -n 2 my_script.py
+
+The ``mlx.launch`` command connects to the provided host and launches the input
+script on each host. It monitors each of the launched processes and terminates
+the rest if one of them fails unexpectedly or if ``mlx.launch`` is terminated.
+It also takes care of forwarding the output of each remote process to stdout
+and stderr respectively.
+
+Providing Hosts
+^^^^^^^^^^^^^^^^
+
+Hosts can be provided as command line arguments, like above, but the way that
+allows to fully define a list of hosts is via a JSON hostfile. The hostfile has
+a very simple schema. It is simply a list of objects that define each host via
+a hostname to ssh to and a list of IPs to utilize for the communication.
+
+.. code:: json
+
+    [
+        {"ssh": "hostname1", "ips": ["123.123.1.1", "123.123.2.1"]},
+        {"ssh": "hostname2", "ips": ["123.123.1.2", "123.123.2.2"]}
+    ]
+
+You can use ``mlx.distributed_config --over ethernet`` to create a hostfile
+with IPs corresponding to the ``en0`` interface.
+
+Setting up Remote Hosts
+^^^^^^^^^^^^^^^^^^^^^^^^
+
+In order to be able to launch the script on each host we need to be able to
+connect via ssh. Moreover the input script and python binary need to be on each
+host and on the same path. A good checklist to debug errors is the following:
+
+* ``ssh hostname`` works without asking for password or host confirmation
+* the python binary is available on all hosts at the same path. You can use
+  ``mlx.launch --print-python`` to see what that path is.
+* the script you want to run is available on all hosts at the same path
+
+.. _mpi_specifics:
+
+MPI Specifics
+-------------
+
+One can use MPI by passing ``--backend mpi`` to ``mlx.launch``. In that case,
+``mlx.launch`` is a thin wrapper over ``mpirun``. Moreover,
+
+* The IPs in the hostfile are ignored
+* The ssh connectivity requirement is stronger as every node needs to be able
+  to connect to every other node
+* ``mpirun`` needs to be available on every node at the same path
+
+Finally, one can pass arguments to ``mpirun`` using ``--mpi-arg``. For instance
+to choose a specific interface for the byte-transfer-layer of MPI we can call
+``mlx.launch`` as follows:
+
+.. code:: shell
+
+    mlx.launch --backend mpi --mpi-arg '--mca btl_tcp_if_include en0' --hostfile hosts.json my_script.py
+
+
+.. _ring_specifics:
+
+Ring Specifics
+--------------
+
+The ring backend, which is also the default backend, can be explicitly selected
+with the argument ``--backend ring``. The ring backend has some specific
+requirements and arguments that are different to MPI:
+
+* The argument ``--hosts`` only accepts IPs and not hostnames. If we need to
+  ssh to a hostname that does not correspond to the IP we want to bind to we
+  have to provide a hostfile.
+* ``--starting-port`` defines the port to bind to on the remote hosts.
+  Specifically rank 0 for the first IP will use this port and each subsequent
+  IP or rank will add 1 to this port.
+* ``--connections-per-ip`` allows us to increase the number of connections
+  between neighboring nodes. This corresponds to ``--mca btl_tcp_links 2`` for
+  ``mpirun``.

mlx/CMakeLists.txt

@@ -17,6 +17,7 @@ target_sources(
           ${CMAKE_CURRENT_SOURCE_DIR}/transforms.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/utils.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/linalg.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/version.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/backend/metal/metal.h)
 
 if(MSVC)

mlx/backend/common/slicing.cpp

@@ -14,6 +14,10 @@ std::tuple<int64_t, Strides> prepare_slice(
     data_offset += start_indices[i] * in.strides()[i];
     inp_strides[i] = in.strides()[i] * strides[i];
   }
+  // Normalize the offset
+  if (data_offset < 0) {
+    data_offset += in.data_size();
+  }
   return std::make_tuple(data_offset, inp_strides);
 }
 
@@ -54,9 +58,10 @@ void slice(
       data_end += end_idx * in.strides()[i];
     }
   }
-  // data_end can be -1
-  size_t data_size =
-      data_end < 0 ? (data_offset - data_end) : (data_end - data_offset);
+  if (data_end < 0) {
+    data_end += in.data_size();
+  }
+  size_t data_size = (data_end - data_offset);
   shared_buffer_slice(in, inp_strides, data_offset, data_size, out);
 }
 

mlx/backend/cpu/cholesky.cpp

@@ -8,6 +8,7 @@
 
 namespace mlx::core {
 
+template <typename T>
 void cholesky_impl(const array& a, array& factor, bool upper) {
   // Lapack uses the column-major convention. We take advantage of the fact that
   // the matrix should be symmetric:
@@ -28,13 +29,12 @@ void cholesky_impl(const array& a, array& factor, bool upper) {
   const int N = a.shape(-1);
   const size_t num_matrices = a.size() / (N * N);
 
-  float* matrix = factor.data<float>();
+  T* matrix = factor.data<T>();
 
   for (int i = 0; i < num_matrices; i++) {
     // Compute Cholesky factorization.
     int info;
-    MLX_LAPACK_FUNC(spotrf)
-    (
+    potrf<T>(
         /* uplo = */ &uplo,
         /* n = */ &N,
         /* a = */ matrix,
@@ -65,10 +65,17 @@ void cholesky_impl(const array& a, array& factor, bool upper) {
 }
 
 void Cholesky::eval_cpu(const std::vector<array>& inputs, array& output) {
-  if (inputs[0].dtype() != float32) {
-    throw std::runtime_error("[Cholesky::eval] only supports float32.");
+  switch (inputs[0].dtype()) {
+    case float32:
+      cholesky_impl<float>(inputs[0], output, upper_);
+      break;
+    case float64:
+      cholesky_impl<double>(inputs[0], output, upper_);
+      break;
+    default:
+      throw std::runtime_error(
+          "[Cholesky::eval_cpu] only supports float32 or float64.");
   }
-  cholesky_impl(inputs[0], output, upper_);
 }
 
 } // namespace mlx::core

mlx/backend/cpu/eigh.cpp

@@ -11,35 +11,64 @@ namespace mlx::core {
 
 namespace {
 
-void ssyevd(
-    char jobz,
-    char uplo,
-    float* a,
-    int N,
-    float* w,
-    float* work,
-    int lwork,
-    int* iwork,
-    int liwork) {
+template <typename T>
+void eigh_impl(
+    array& vectors,
+    array& values,
+    const std::string& uplo,
+    bool compute_eigenvectors) {
+  auto vec_ptr = vectors.data<T>();
+  auto eig_ptr = values.data<T>();
+
+  char jobz = compute_eigenvectors ? 'V' : 'N';
+  auto N = vectors.shape(-1);
+
+  // Work query
+  int lwork = -1;
+  int liwork = -1;
   int info;
-  MLX_LAPACK_FUNC(ssyevd)
-  (
-      /* jobz = */ &jobz,
-      /* uplo = */ &uplo,
-      /* n = */ &N,
-      /* a = */ a,
-      /* lda = */ &N,
-      /* w = */ w,
-      /* work = */ work,
-      /* lwork = */ &lwork,
-      /* iwork = */ iwork,
-      /* liwork = */ &liwork,
-      /* info = */ &info);
-  if (info != 0) {
-    std::stringstream msg;
-    msg << "[Eigh::eval_cpu] Eigenvalue decomposition failed with error code "
-        << info;
-    throw std::runtime_error(msg.str());
+  {
+    T work;
+    int iwork;
+    syevd<T>(
+        &jobz,
+        uplo.c_str(),
+        &N,
+        nullptr,
+        &N,
+        nullptr,
+        &work,
+        &lwork,
+        &iwork,
+        &liwork,
+        &info);
+    lwork = static_cast<int>(work);
+    liwork = iwork;
+  }
+
+  auto work_buf = array::Data{allocator::malloc_or_wait(sizeof(T) * lwork)};
+  auto iwork_buf = array::Data{allocator::malloc_or_wait(sizeof(int) * liwork)};
+  for (size_t i = 0; i < vectors.size() / (N * N); ++i) {
+    syevd<T>(
+        &jobz,
+        uplo.c_str(),
+        &N,
+        vec_ptr,
+        &N,
+        eig_ptr,
+        static_cast<T*>(work_buf.buffer.raw_ptr()),
+        &lwork,
+        static_cast<int*>(iwork_buf.buffer.raw_ptr()),
+        &liwork,
+        &info);
+    vec_ptr += N * N;
+    eig_ptr += N;
+    if (info != 0) {
+      std::stringstream msg;
+      msg << "[Eigh::eval_cpu] Eigenvalue decomposition failed with error code "
+          << info;
+      throw std::runtime_error(msg.str());
+    }
   }
 }
 
@@ -80,39 +109,16 @@ void Eigh::eval_cpu(
     }
     vectors.move_shared_buffer(vectors, strides, flags, vectors.data_size());
   }
-
-  auto vec_ptr = vectors.data<float>();
-  auto eig_ptr = values.data<float>();
-
-  char jobz = compute_eigenvectors_ ? 'V' : 'N';
-  auto N = a.shape(-1);
-
-  // Work query
-  int lwork;
-  int liwork;
-  {
-    float work;
-    int iwork;
-    ssyevd(jobz, uplo_[0], nullptr, N, nullptr, &work, -1, &iwork, -1);
-    lwork = static_cast<int>(work);
-    liwork = iwork;
-  }
-
-  auto work_buf = array::Data{allocator::malloc_or_wait(sizeof(float) * lwork)};
-  auto iwork_buf = array::Data{allocator::malloc_or_wait(sizeof(int) * liwork)};
-  for (size_t i = 0; i < a.size() / (N * N); ++i) {
-    ssyevd(
-        jobz,
-        uplo_[0],
-        vec_ptr,
-        N,
-        eig_ptr,
-        static_cast<float*>(work_buf.buffer.raw_ptr()),
-        lwork,
-        static_cast<int*>(iwork_buf.buffer.raw_ptr()),
-        liwork);
-    vec_ptr += N * N;
-    eig_ptr += N;
+  switch (a.dtype()) {
+    case float32:
+      eigh_impl<float>(vectors, values, uplo_, compute_eigenvectors_);
+      break;
+    case float64:
+      eigh_impl<double>(vectors, values, uplo_, compute_eigenvectors_);
+      break;
+    default:
+      throw std::runtime_error(
+          "[Eigh::eval_cpu] only supports float32 or float64.");
   }
 }
 

mlx/backend/cpu/gemms/bnns.cpp

@@ -43,6 +43,8 @@ void matmul_bnns(
 
   BNNSDataType bnns_dtype = to_bnns_dtype(out.dtype());
 
+#pragma GCC diagnostic push
+#pragma GCC diagnostic ignored "-Wdeprecated-declarations"
   const BNNSLayerParametersBroadcastMatMul gemm_params{
       /* float alpha = */ alpha,
       /* float beta = */ beta,
@@ -124,6 +126,7 @@ void matmul_bnns(
   }
 
   BNNSFilterDestroy(bnns_filter);
+#pragma GCC diagnostic pop
 }
 
 template <>

mlx/backend/cpu/inverse.cpp

@@ -5,44 +5,33 @@
 #include "mlx/backend/cpu/lapack.h"
 #include "mlx/primitives.h"
 
-int strtri_wrapper(char uplo, char diag, float* matrix, int N) {
-  int info;
-  MLX_LAPACK_FUNC(strtri)
-  (
-      /* uplo = */ &uplo,
-      /* diag = */ &diag,
-      /* N = */ &N,
-      /* a = */ matrix,
-      /* lda = */ &N,
-      /* info = */ &info);
-  return info;
-}
-
 namespace mlx::core {
 
+template <typename T>
 void general_inv(array& inv, int N, int i) {
   int info;
   auto ipiv = array::Data{allocator::malloc_or_wait(sizeof(int) * N)};
   // Compute LU factorization.
-  sgetrf_(
+  getrf<T>(
       /* m = */ &N,
       /* n = */ &N,
-      /* a = */ inv.data<float>() + N * N * i,
+      /* a = */ inv.data<T>() + N * N * i,
       /* lda = */ &N,
       /* ipiv = */ static_cast<int*>(ipiv.buffer.raw_ptr()),
       /* info = */ &info);
 
   if (info != 0) {
     std::stringstream ss;
-    ss << "inverse_impl: LU factorization failed with error code " << info;
+    ss << "[Inverse::eval_cpu] LU factorization failed with error code "
+       << info;
     throw std::runtime_error(ss.str());
   }
 
   static const int lwork_query = -1;
-  float workspace_size = 0;
+  T workspace_size = 0;
 
   // Compute workspace size.
-  sgetri_(
+  getri<T>(
       /* m = */ &N,
       /* a = */ nullptr,
       /* lda = */ &N,
@@ -53,42 +42,67 @@ void general_inv(array& inv, int N, int i) {
 
   if (info != 0) {
     std::stringstream ss;
-    ss << "inverse_impl: LU workspace calculation failed with error code "
+    ss << "[Inverse::eval_cpu] LU workspace calculation failed with error code "
        << info;
     throw std::runtime_error(ss.str());
   }
 
   const int lwork = workspace_size;
-  auto scratch = array::Data{allocator::malloc_or_wait(sizeof(float) * lwork)};
+  auto scratch = array::Data{allocator::malloc_or_wait(sizeof(T) * lwork)};
 
   // Compute inverse.
-  sgetri_(
+  getri<T>(
       /* m = */ &N,
-      /* a = */ inv.data<float>() + N * N * i,
+      /* a = */ inv.data<T>() + N * N * i,
       /* lda = */ &N,
       /* ipiv = */ static_cast<int*>(ipiv.buffer.raw_ptr()),
-      /* work = */ static_cast<float*>(scratch.buffer.raw_ptr()),
+      /* work = */ static_cast<T*>(scratch.buffer.raw_ptr()),
       /* lwork = */ &lwork,
       /* info = */ &info);
 
   if (info != 0) {
     std::stringstream ss;
-    ss << "inverse_impl: inversion failed with error code " << info;
+    ss << "[Inverse::eval_cpu] inversion failed with error code " << info;
     throw std::runtime_error(ss.str());
   }
 }
 
+template <typename T>
 void tri_inv(array& inv, int N, int i, bool upper) {
   const char uplo = upper ? 'L' : 'U';
   const char diag = 'N';
-  int info = strtri_wrapper(uplo, diag, inv.data<float>() + N * N * i, N);
+  T* data = inv.data<T>() + N * N * i;
+  int info;
+  trtri<T>(
+      /* uplo = */ &uplo,
+      /* diag = */ &diag,
+      /* N = */ &N,
+      /* a = */ data,
+      /* lda = */ &N,
+      /* info = */ &info);
+
+  // zero out the other triangle
+  if (upper) {
+    for (int i = 0; i < N; i++) {
+      std::fill(data, data + i, 0.0f);
+      data += N;
+    }
+  } else {
+    for (int i = 0; i < N; i++) {
+      std::fill(data + i + 1, data + N, 0.0f);
+      data += N;
+    }
+  }
+
   if (info != 0) {
     std::stringstream ss;
-    ss << "inverse_impl: triangular inversion failed with error code " << info;
+    ss << "[Inverse::eval_cpu] triangular inversion failed with error code "
+       << info;
     throw std::runtime_error(ss.str());
   }
 }
 
+template <typename T>
 void inverse_impl(const array& a, array& inv, bool tri, bool upper) {
   // Lapack uses the column-major convention. We take advantage of the following
   // identity to avoid transposing (see
@@ -103,18 +117,25 @@ void inverse_impl(const array& a, array& inv, bool tri, bool upper) {
 
   for (int i = 0; i < num_matrices; i++) {
     if (tri) {
-      tri_inv(inv, N, i, upper);
+      tri_inv<T>(inv, N, i, upper);
     } else {
-      general_inv(inv, N, i);
+      general_inv<T>(inv, N, i);
     }
   }
 }
 
 void Inverse::eval_cpu(const std::vector<array>& inputs, array& output) {
-  if (inputs[0].dtype() != float32) {
-    throw std::runtime_error("[Inverse::eval] only supports float32.");
+  switch (inputs[0].dtype()) {
+    case float32:
+      inverse_impl<float>(inputs[0], output, tri_, upper_);
+      break;
+    case float64:
+      inverse_impl<double>(inputs[0], output, tri_, upper_);
+      break;
+    default:
+      throw std::runtime_error(
+          "[Inverse::eval_cpu] only supports float32 or float64.");
   }
-  inverse_impl(inputs[0], output, tri_, upper_);
 }
 
 } // namespace mlx::core

mlx/backend/cpu/lapack.h

@@ -31,3 +31,22 @@
 #define MLX_LAPACK_FUNC(f) f##_
 
 #endif
+
+#define INSTANTIATE_LAPACK_TYPES(FUNC)                       \
+  template <typename T, typename... Args>                    \
+  void FUNC(Args... args) {                                  \
+    if constexpr (std::is_same_v<T, float>) {                \
+      MLX_LAPACK_FUNC(s##FUNC)(std::forward<Args>(args)...); \
+    } else if constexpr (std::is_same_v<T, double>) {        \
+      MLX_LAPACK_FUNC(d##FUNC)(std::forward<Args>(args)...); \
+    }                                                        \
+  }
+
+INSTANTIATE_LAPACK_TYPES(geqrf)
+INSTANTIATE_LAPACK_TYPES(orgqr)
+INSTANTIATE_LAPACK_TYPES(syevd)
+INSTANTIATE_LAPACK_TYPES(potrf)
+INSTANTIATE_LAPACK_TYPES(gesvdx)
+INSTANTIATE_LAPACK_TYPES(getrf)
+INSTANTIATE_LAPACK_TYPES(getri)
+INSTANTIATE_LAPACK_TYPES(trtri)

mlx/backend/cpu/luf.cpp

@@ -9,11 +9,8 @@
 
 namespace mlx::core {
 
-void lu_factor_impl(
-    const array& a,
-    array& lu,
-    array& pivots,
-    array& row_indices) {
+template <typename T>
+void luf_impl(const array& a, array& lu, array& pivots, array& row_indices) {
   int M = a.shape(-2);
   int N = a.shape(-1);
 
@@ -31,7 +28,7 @@ void lu_factor_impl(
   copy_inplace(
       a, lu, a.shape(), a.strides(), strides, 0, 0, CopyType::GeneralGeneral);
 
-  auto a_ptr = lu.data<float>();
+  auto a_ptr = lu.data<T>();
 
   pivots.set_data(allocator::malloc_or_wait(pivots.nbytes()));
   row_indices.set_data(allocator::malloc_or_wait(row_indices.nbytes()));
@@ -42,13 +39,13 @@ void lu_factor_impl(
   size_t num_matrices = a.size() / (M * N);
   for (size_t i = 0; i < num_matrices; ++i) {
     // Compute LU factorization of A
-    MLX_LAPACK_FUNC(sgetrf)
-    (/* m */ &M,
-     /* n */ &N,
-     /* a */ a_ptr,
-     /* lda */ &M,
-     /* ipiv */ reinterpret_cast<int*>(pivots_ptr),
-     /* info */ &info);
+    getrf<T>(
+        /* m */ &M,
+        /* n */ &N,
+        /* a */ a_ptr,
+        /* lda */ &M,
+        /* ipiv */ reinterpret_cast<int*>(pivots_ptr),
+        /* info */ &info);
 
     if (info != 0) {
       std::stringstream ss;
@@ -59,10 +56,14 @@ void lu_factor_impl(
     }
 
     // Subtract 1 to get 0-based index
-    for (int j = 0; j < pivots.shape(-1); ++j) {
+    int j = 0;
+    for (; j < pivots.shape(-1); ++j) {
       pivots_ptr[j]--;
       row_indices_ptr[j] = j;
     }
+    for (; j < row_indices.shape(-1); ++j) {
+      row_indices_ptr[j] = j;
+    }
     for (int j = pivots.shape(-1) - 1; j >= 0; --j) {
       auto piv = pivots_ptr[j];
       auto t1 = row_indices_ptr[piv];
@@ -82,7 +83,17 @@ void LUF::eval_cpu(
     const std::vector<array>& inputs,
     std::vector<array>& outputs) {
   assert(inputs.size() == 1);
-  lu_factor_impl(inputs[0], outputs[0], outputs[1], outputs[2]);
+  switch (inputs[0].dtype()) {
+    case float32:
+      luf_impl<float>(inputs[0], outputs[0], outputs[1], outputs[2]);
+      break;
+    case float64:
+      luf_impl<double>(inputs[0], outputs[0], outputs[1], outputs[2]);
+      break;
+    default:
+      throw std::runtime_error(
+          "[LUF::eval_cpu] only supports float32 or float64.");
+  }
 }
 
 } // namespace mlx::core

mlx/backend/cpu/make_compiled_preamble.sh

@@ -18,13 +18,16 @@ if [ "$CLANG" = "TRUE" ]; then
 #include <complex>
 #include <cstdint>
 #include <vector>
+#ifdef __ARM_FEATURE_FP16_SCALAR_ARITHMETIC
+#include <arm_fp16.h>
+#endif
 EOM
-CC_FLAGS="-arch ${ARCH}"
+CC_FLAGS="-arch ${ARCH} -nobuiltininc -nostdinc"
 else
 CC_FLAGS="-std=c++17"
 fi
 
-CONTENT=$($GCC $CC_FLAGS -I "$SRCDIR" -E "$SRCDIR/mlx/backend/cpu/compiled_preamble.h" 2>/dev/null)
+CONTENT=$($GCC $CC_FLAGS -I "$SRCDIR" -E -P "$SRCDIR/mlx/backend/cpu/compiled_preamble.h" 2>/dev/null)
 
 cat << EOF > "$OUTPUT_FILE"
 const char* get_kernel_preamble() {

mlx/backend/cpu/qrf.cpp

@@ -7,36 +7,6 @@
 
 namespace mlx::core {
 
-template <typename T>
-struct lpack;
-
-template <>
-struct lpack<float> {
-  static void xgeqrf(
-      const int* m,
-      const int* n,
-      float* a,
-      const int* lda,
-      float* tau,
-      float* work,
-      const int* lwork,
-      int* info) {
-    sgeqrf_(m, n, a, lda, tau, work, lwork, info);
-  }
-  static void xorgqr(
-      const int* m,
-      const int* n,
-      const int* k,
-      float* a,
-      const int* lda,
-      const float* tau,
-      float* work,
-      const int* lwork,
-      int* info) {
-    sorgqr_(m, n, k, a, lda, tau, work, lwork, info);
-  }
-};
-
 template <typename T>
 void qrf_impl(const array& a, array& q, array& r) {
   const int M = a.shape(-2);
@@ -48,7 +18,7 @@ void qrf_impl(const array& a, array& q, array& r) {
       allocator::malloc_or_wait(sizeof(T) * num_matrices * num_reflectors);
 
   // Copy A to inplace input and make it col-contiguous
-  array in(a.shape(), float32, nullptr, {});
+  array in(a.shape(), a.dtype(), nullptr, {});
   auto flags = in.flags();
 
   // Copy the input to be column contiguous
@@ -66,8 +36,7 @@ void qrf_impl(const array& a, array& q, array& r) {
   int info;
 
   // Compute workspace size
-  lpack<T>::xgeqrf(
-      &M, &N, nullptr, &lda, nullptr, &optimal_work, &lwork, &info);
+  geqrf<T>(&M, &N, nullptr, &lda, nullptr, &optimal_work, &lwork, &info);
 
   // Update workspace size
   lwork = optimal_work;
@@ -76,10 +45,10 @@ void qrf_impl(const array& a, array& q, array& r) {
   // Loop over matrices
   for (int i = 0; i < num_matrices; ++i) {
     // Solve
-    lpack<T>::xgeqrf(
+    geqrf<T>(
         &M,
         &N,
-        in.data<float>() + M * N * i,
+        in.data<T>() + M * N * i,
         &lda,
         static_cast<T*>(tau.raw_ptr()) + num_reflectors * i,
         static_cast<T*>(work.raw_ptr()),
@@ -105,7 +74,7 @@ void qrf_impl(const array& a, array& q, array& r) {
 
   // Get work size
   lwork = -1;
-  lpack<T>::xorgqr(
+  orgqr<T>(
       &M,
       &num_reflectors,
       &num_reflectors,
@@ -121,11 +90,11 @@ void qrf_impl(const array& a, array& q, array& r) {
   // Loop over matrices
   for (int i = 0; i < num_matrices; ++i) {
     // Compute Q
-    lpack<T>::xorgqr(
+    orgqr<T>(
         &M,
         &num_reflectors,
         &num_reflectors,
-        in.data<float>() + M * N * i,
+        in.data<T>() + M * N * i,
         &lda,
         static_cast<T*>(tau.raw_ptr()) + num_reflectors * i,
         static_cast<T*>(work.raw_ptr()),
@@ -152,10 +121,17 @@ void qrf_impl(const array& a, array& q, array& r) {
 void QRF::eval_cpu(
     const std::vector<array>& inputs,
     std::vector<array>& outputs) {
-  if (!(inputs[0].dtype() == float32)) {
-    throw std::runtime_error("[QRF::eval] only supports float32.");
+  switch (inputs[0].dtype()) {
+    case float32:
+      qrf_impl<float>(inputs[0], outputs[0], outputs[1]);
+      break;
+    case float64:
+      qrf_impl<double>(inputs[0], outputs[0], outputs[1]);
+      break;
+    default:
+      throw std::runtime_error(
+          "[QRF::eval_cpu] only supports float32 or float64.");
   }
-  qrf_impl<float>(inputs[0], outputs[0], outputs[1]);
 }
 
 } // namespace mlx::core

mlx/backend/cpu/quantized.cpp

@@ -164,8 +164,8 @@ simd::Simd<uint32_t, S> extract_bits_simd(const uint32_t* w) {
   } else if constexpr (bits == 8 && S == 8) {
     constexpr std::array<uint32_t, 8> shifts_ = {{0, 8, 16, 24, 0, 8, 16, 24}};
     auto shifts(*(simd::Simd<uint32_t, S>*)&shifts_);
-    auto l = simd::Simd<uint32_t, 4>(*w++);
-    auto r = simd::Simd<uint32_t, 4>(*w);
+    auto l = simd::Simd<uint32_t, S / 2>(*w++);
+    auto r = simd::Simd<uint32_t, S / 2>(*w);
     wi = simd::Simd<uint32_t, S>(l, r);
     wi = wi >> shifts;
     wi = wi & bitmask;
@@ -543,8 +543,8 @@ void quantize(
   T* scales = scales_.data<T>();
   T* biases = biases_.data<T>();
 
-  T n_bins = (1 << bits) - 1;
-  T eps = 1e-7;
+  float n_bins = (1 << bits) - 1;
+  float eps = 1e-7;
   bool power_of_2_bits = is_power_of_2(bits);
   int el_per_int = bits == 3 ? 8 : bits == 6 ? 4 : 32 / bits;
   // For 3/6 bits we read 3 uint8s at a time instead of 1 uint32
@@ -554,32 +554,30 @@ void quantize(
 
   for (size_t i = 0; i < n_groups; ++i) {
     size_t w_idx = i * group_size;
-    T w_min = std::numeric_limits<float>::infinity();
-    T w_max = -w_min;
+    float w_min = std::numeric_limits<float>::infinity();
+    float w_max = -w_min;
     for (int j = 0; j < group_size; ++j) {
-      w_max = std::max(w_max, w[w_idx + j]);
-      w_min = std::min(w_min, w[w_idx + j]);
+      w_max = std::max(w_max, (float)w[w_idx + j]);
+      w_min = std::min(w_min, (float)w[w_idx + j]);
     }
     bool mask = std::abs(w_min) > std::abs(w_max);
-    T scale = std::max(T((w_max - w_min) / n_bins), eps);
+    float scale = std::max((w_max - w_min) / n_bins, eps);
     scale = mask ? scale : -scale;
 
-    auto edge = mask ? w_min : w_max;
-    auto q0 = std::rint(edge / scale);
-    if (q0 == 0) {
-      scales[i] = scale;
-      biases[i] = 0;
-    } else {
-      scales[i] = edge / q0;
-      biases[i] = edge;
+    float edge = mask ? w_min : w_max;
+    float q0 = std::rint(edge / scale);
+    float bias = 0;
+    if (q0 != 0) {
+      scale = edge / q0;
+      bias = edge;
     }
     size_t out_idx = i * int_per_group;
     for (int j = 0; j < int_per_group / bytes_per_pack; ++j) {
       uint32_t out_el = 0;
       for (int k = 0; k < el_per_int; ++k) {
-        T w_el = w[w_idx + j * el_per_int + k];
-        w_el = std::rint((w_el - biases[i]) / scales[i]);
-        w_el = std::min(std::max(w_el, T(0)), n_bins);
+        float w_el = w[w_idx + j * el_per_int + k];
+        w_el = std::rint((w_el - bias) / scale);
+        w_el = std::min(std::max(w_el, 0.0f), n_bins);
         out_el |= static_cast<uint32_t>(w_el) << (k * bits);
       }
       if (power_of_2_bits) {
@@ -590,6 +588,8 @@ void quantize(
         out[out_idx + bytes_per_pack * j + 2] = (out_el & 0xff0000) >> 16;
       }
     }
+    scales[i] = static_cast<T>(scale);
+    biases[i] = static_cast<T>(bias);
   }
 }
 

mlx/backend/cpu/simd/math.h

@@ -186,7 +186,7 @@ Simd<T, N> erfinv(Simd<T, N> a_) {
       return a * rhs(t);
     }
   } else {
-    return a * select(t > thresh, lhs(t), rhs(t));
+    return a * select(abs(t) > thresh, lhs(t), rhs(t));
   }
 }
 

mlx/backend/cpu/svd.cpp

@@ -7,7 +7,8 @@
 
 namespace mlx::core {
 
-void svd_impl(const array& a, array& u, array& s, array& vt) {
+template <typename T>
+void svd_impl(const array& a, T* u_data, T* s_data, T* vt_data) {
   // Lapack uses the column-major convention. To avoid having to transpose
   // the input and then transpose the outputs, we swap the indices/sizes of the
   // matrices and take advantage of the following identity (see
@@ -31,21 +32,16 @@ void svd_impl(const array& a, array& u, array& s, array& vt) {
   size_t num_matrices = a.size() / (M * N);
 
   // lapack clobbers the input, so we have to make a copy.
-  array in(a.shape(), float32, nullptr, {});
+  array in(a.shape(), a.dtype(), nullptr, {});
   copy(a, in, a.flags().row_contiguous ? CopyType::Vector : CopyType::General);
 
-  // Allocate outputs.
-  u.set_data(allocator::malloc_or_wait(u.nbytes()));
-  s.set_data(allocator::malloc_or_wait(s.nbytes()));
-  vt.set_data(allocator::malloc_or_wait(vt.nbytes()));
-
-  static constexpr auto job_u = "V";
-  static constexpr auto job_vt = "V";
+  auto job_u = (u_data && vt_data) ? "V" : "N";
+  auto job_vt = (u_data && vt_data) ? "V" : "N";
   static constexpr auto range = "A";
 
   // Will contain the number of singular values after the call has returned.
   int ns = 0;
-  float workspace_dimension = 0;
+  T workspace_dimension = 0;
 
   // Will contain the indices of eigenvectors that failed to converge (not used
   // here but required by lapack).
@@ -54,13 +50,13 @@ void svd_impl(const array& a, array& u, array& s, array& vt) {
   static const int lwork_query = -1;
 
   static const int ignored_int = 0;
-  static const float ignored_float = 0;
+  static const T ignored_float = 0;
+  static T ignored_output = 0;
 
   int info;
 
   // Compute workspace size.
-  MLX_LAPACK_FUNC(sgesvdx)
-  (
+  gesvdx<T>(
       /* jobu = */ job_u,
       /* jobvt = */ job_vt,
       /* range = */ range,
@@ -86,64 +82,91 @@ void svd_impl(const array& a, array& u, array& s, array& vt) {
 
   if (info != 0) {
     std::stringstream ss;
-    ss << "svd_impl: sgesvdx_ workspace calculation failed with code " << info;
+    ss << "[SVD::eval_cpu] workspace calculation failed with code " << info;
     throw std::runtime_error(ss.str());
   }
 
   const int lwork = workspace_dimension;
-  auto scratch = array::Data{allocator::malloc_or_wait(sizeof(float) * lwork)};
+  auto scratch = array::Data{allocator::malloc_or_wait(sizeof(T) * lwork)};
 
   // Loop over matrices.
   for (int i = 0; i < num_matrices; i++) {
-    MLX_LAPACK_FUNC(sgesvdx)
-    (
+    gesvdx<T>(
         /* jobu = */ job_u,
         /* jobvt = */ job_vt,
         /* range = */ range,
         // M and N are swapped since lapack expects column-major.
         /* m = */ &N,
         /* n = */ &M,
-        /* a = */ in.data<float>() + M * N * i,
+        /* a = */ in.data<T>() + M * N * i,
         /* lda = */ &lda,
         /* vl = */ &ignored_float,
         /* vu = */ &ignored_float,
         /* il = */ &ignored_int,
         /* iu = */ &ignored_int,
         /* ns = */ &ns,
-        /* s = */ s.data<float>() + K * i,
+        /* s = */ s_data + K * i,
         // According to the identity above, lapack will write Vᵀᵀ as U.
-        /* u = */ vt.data<float>() + N * N * i,
+        /* u = */ vt_data ? vt_data + N * N * i : &ignored_output,
         /* ldu = */ &ldu,
         // According to the identity above, lapack will write Uᵀ as Vᵀ.
-        /* vt = */ u.data<float>() + M * M * i,
+        /* vt = */ u_data ? u_data + M * M * i : &ignored_output,
         /* ldvt = */ &ldvt,
-        /* work = */ static_cast<float*>(scratch.buffer.raw_ptr()),
+        /* work = */ static_cast<T*>(scratch.buffer.raw_ptr()),
         /* lwork = */ &lwork,
         /* iwork = */ static_cast<int*>(iwork.buffer.raw_ptr()),
         /* info = */ &info);
 
     if (info != 0) {
       std::stringstream ss;
-      ss << "svd_impl: sgesvdx_ failed with code " << info;
+      ss << "[SVD::eval_cpu] failed with code " << info;
       throw std::runtime_error(ss.str());
     }
 
     if (ns != K) {
       std::stringstream ss;
-      ss << "svd_impl: expected " << K << " singular values, but " << ns
+      ss << "[SVD::eval_cpu] expected " << K << " singular values, but " << ns
          << " were computed.";
       throw std::runtime_error(ss.str());
     }
   }
 }
 
+template <typename T>
+void compute_svd(const array& a, bool compute_uv, std::vector<array>& outputs) {
+  if (compute_uv) {
+    array& u = outputs[0];
+    array& s = outputs[1];
+    array& vt = outputs[2];
+
+    u.set_data(allocator::malloc_or_wait(u.nbytes()));
+    s.set_data(allocator::malloc_or_wait(s.nbytes()));
+    vt.set_data(allocator::malloc_or_wait(vt.nbytes()));
+
+    svd_impl<T>(a, u.data<T>(), s.data<T>(), vt.data<T>());
+  } else {
+    array& s = outputs[0];
+
+    s.set_data(allocator::malloc_or_wait(s.nbytes()));
+
+    svd_impl<T>(a, nullptr, s.data<T>(), nullptr);
+  }
+}
+
 void SVD::eval_cpu(
     const std::vector<array>& inputs,
     std::vector<array>& outputs) {
-  if (!(inputs[0].dtype() == float32)) {
-    throw std::runtime_error("[SVD::eval] only supports float32.");
+  switch (inputs[0].dtype()) {
+    case float32:
+      compute_svd<float>(inputs[0], compute_uv_, outputs);
+      break;
+    case float64:
+      compute_svd<double>(inputs[0], compute_uv_, outputs);
+      break;
+    default:
+      throw std::runtime_error(
+          "[SVD::eval_cpu] only supports float32 or float64.");
   }
-  svd_impl(inputs[0], outputs[0], outputs[1], outputs[2]);
 }
 
 } // namespace mlx::core

mlx/backend/metal/allocator.cpp

@@ -10,6 +10,9 @@
 
 namespace mlx::core {
 
+constexpr size_t resource_options =
+    MTL::ResourceStorageModeShared | MTL::ResourceHazardTrackingModeUntracked;
+
 namespace allocator {
 
 Allocator& allocator() {
@@ -150,15 +153,34 @@ MetalAllocator::MetalAllocator()
     : device_(device(mlx::core::Device::gpu).mtl_device()),
       residency_set_(device_),
       buffer_cache_(device_) {
-  auto memsize = std::get<size_t>(device_info()["memory_size"]);
+  auto pool = metal::new_scoped_memory_pool();
+  auto memsize = std::get<size_t>(device_info().at("memory_size"));
   auto max_rec_size =
-      std::get<size_t>(device_info()["max_recommended_working_set_size"]);
-  resource_limit_ = std::get<size_t>(device_info()["resource_limit"]);
+      std::get<size_t>(device_info().at("max_recommended_working_set_size"));
+  resource_limit_ = std::get<size_t>(device_info().at("resource_limit"));
   block_limit_ = std::min(1.5 * max_rec_size, 0.95 * memsize);
   gc_limit_ = std::min(static_cast<size_t>(0.95 * max_rec_size), block_limit_);
   max_pool_size_ = block_limit_;
   device(mlx::core::Device::gpu)
       .set_residency_set(residency_set_.mtl_residency_set());
+  bool is_vm = std::get<std::string>(device_info().at("device_name")) ==
+      "Apple Paravirtual device";
+  if (is_vm) {
+    return;
+  }
+  auto heap_desc = MTL::HeapDescriptor::alloc()->init();
+  heap_desc->setResourceOptions(resource_options);
+  heap_desc->setSize(heap_size_);
+  heap_ = device_->newHeap(heap_desc);
+  heap_desc->release();
+  residency_set_.insert(heap_);
+}
+
+MetalAllocator::~MetalAllocator() {
+  auto pool = metal::new_scoped_memory_pool();
+  if (heap_) {
+    heap_->release();
+  }
 }
 
 size_t MetalAllocator::set_cache_limit(size_t limit) {
@@ -226,8 +248,6 @@ Buffer MetalAllocator::malloc(size_t size, bool allow_swap /* = false */) {
     }
 
     // Allocate new buffer if needed
-    size_t res_opt = MTL::ResourceStorageModeShared;
-    res_opt |= MTL::ResourceHazardTrackingModeUntracked;
     if (num_resources_ >= resource_limit_) {
       std::ostringstream msg;
       msg << "[metal::malloc] Resource limit (" << resource_limit_
@@ -235,7 +255,12 @@ Buffer MetalAllocator::malloc(size_t size, bool allow_swap /* = false */) {
       throw std::runtime_error(msg.str());
     }
     lk.unlock();
-    buf = device_->newBuffer(size, res_opt);
+    if (size < small_size_ && heap_) {
+      buf = heap_->newBuffer(size, resource_options);
+    }
+    if (!buf) {
+      buf = device_->newBuffer(size, resource_options);
+    }
     lk.lock();
     if (buf) {
       num_resources_++;
@@ -246,13 +271,15 @@ Buffer MetalAllocator::malloc(size_t size, bool allow_swap /* = false */) {
   peak_memory_ = std::max(peak_memory_, active_memory_);
 
   // Maintain the cache below the requested limit
-  if (get_cache_memory() >= max_pool_size_) {
+  if (get_cache_memory() > max_pool_size_) {
     auto pool = metal::new_scoped_memory_pool();
     num_resources_ -= buffer_cache_.release_cached_buffers(
         get_cache_memory() - max_pool_size_);
   }
 
-  residency_set_.insert(buf);
+  if (!buf->heap()) {
+    residency_set_.insert(buf);
+  }
 
   return Buffer{static_cast<void*>(buf)};
 }
@@ -269,7 +296,9 @@ void MetalAllocator::free(Buffer buffer) {
     return;
   }
   std::unique_lock lk(mutex_);
-  residency_set_.erase(buf);
+  if (!buf->heap()) {
+    residency_set_.erase(buf);
+  }
   active_memory_ -= buf->length();
   if (get_cache_memory() < max_pool_size_) {
     buffer_cache_.recycle_to_cache(buf);
@@ -301,7 +330,7 @@ size_t set_memory_limit(size_t limit, bool relaxed /* = true */) {
 }
 size_t set_wired_limit(size_t limit) {
   if (limit >
-      std::get<size_t>(device_info()["max_recommended_working_set_size"])) {
+      std::get<size_t>(device_info().at("max_recommended_working_set_size"))) {
     throw std::invalid_argument(
         "[metal::set_wired_limit] Setting a wired limit larger than "
         "the maximum working set size is not allowed.");

mlx/backend/metal/allocator.h

@@ -43,6 +43,7 @@ class BufferCache {
   void remove_from_list(BufferHolder* to_remove);
 
   MTL::Device* device_;
+  MTL::Heap* heap_{nullptr};
 
   std::multimap<size_t, BufferHolder*> buffer_pool_;
   BufferHolder* head_;
@@ -78,7 +79,15 @@ class MetalAllocator : public allocator::Allocator {
 
  private:
   MTL::Device* device_;
+
+  // The size of allocations which go on the heap until it is full. This size
+  // is chosen because it is the actual minimum size of a buffer allocated from
+  // the heap, a heap can have at most heap.size() / 256 buffers.
+  static constexpr int small_size_ = 256;
+  static constexpr int heap_size_ = 1 << 20;
+  MTL::Heap* heap_;
   MetalAllocator();
+  ~MetalAllocator();
   friend MetalAllocator& allocator();
 
   // Caching allocator

mlx/backend/metal/conv.cpp

@@ -533,45 +533,6 @@ void implicit_gemm_conv_2D_general_gpu(
   compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
 }
 
-void winograd_conv_2D_fused_gpu(
-    const Stream& s,
-    metal::Device& d,
-    const array& in,
-    const array& wt,
-    array out,
-    const MLXConvParams<2>& conv_params,
-    std::vector<array>& copies_w) {
-  int O_c = conv_params.O;
-  int C_c = conv_params.C;
-
-  int N_tiles_n = conv_params.N;
-  int N_tiles_h = (conv_params.oS[0] + 1) / 2;
-  int N_tiles_w = (conv_params.oS[1] + 1) / 2;
-  int N_tiles = N_tiles_n * N_tiles_h * N_tiles_w;
-
-  int bc = 32;
-  int wm = 4;
-  int wn = 1;
-  std::ostringstream kname;
-  kname << "winograd_conv_2d_fused_" << type_to_name(out) << "_flip"
-        << conv_params.flip;
-  auto& compute_encoder = d.get_command_encoder(s.index);
-  auto kernel = d.get_kernel(kname.str());
-  compute_encoder.set_compute_pipeline_state(kernel);
-
-  compute_encoder.set_input_array(in, 0);
-  compute_encoder.set_input_array(wt, 1);
-  compute_encoder.set_output_array(out, 2);
-
-  compute_encoder.set_bytes(conv_params, 3);
-
-  MTL::Size group_dims = MTL::Size(8, 8, 2);
-  MTL::Size grid_dims =
-      MTL::Size(O_c / 8, (N_tiles_h * N_tiles_w) / 8, N_tiles_n);
-
-  compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
-}
-
 void winograd_conv_2D_gpu(
     const Stream& s,
     metal::Device& d,
@@ -580,6 +541,67 @@ void winograd_conv_2D_gpu(
     array out,
     const MLXConvParams<2>& conv_params,
     std::vector<array>& copies_w) {
+  Shape padded_shape = {
+      conv_params.N,
+      conv_params.iS[0] + 2 * conv_params.pad[0],
+      conv_params.iS[1] + 2 * conv_params.pad[1],
+      conv_params.C};
+
+  padded_shape[1] = 6 * ((padded_shape[1] - 2 + 5) / 6) + 2;
+  padded_shape[2] = 6 * ((padded_shape[2] - 2 + 5) / 6) + 2;
+
+  array in_padded(std::move(padded_shape), in.dtype(), nullptr, {});
+
+  // Fill with zeros
+  array zero_arr = array(0, in.dtype());
+  fill_gpu(zero_arr, in_padded, s);
+  copies_w.push_back(zero_arr);
+
+  // Pick input slice from padded
+  size_t data_offset = conv_params.pad[0] * in_padded.strides()[1] +
+      conv_params.pad[1] * in_padded.strides()[2];
+  array in_padded_slice(in.shape(), in_padded.dtype(), nullptr, {});
+  in_padded_slice.copy_shared_buffer(
+      in_padded,
+      in_padded.strides(),
+      in_padded.flags(),
+      in_padded_slice.size(),
+      data_offset);
+
+  // Copy input values into the slice
+  copy_gpu_inplace(in, in_padded_slice, CopyType::GeneralGeneral, s);
+
+  copies_w.push_back(in_padded_slice);
+  copies_w.push_back(in_padded);
+
+  MLXConvParams<2> conv_params_updated{
+      /* const int  N = */ static_cast<int>(in_padded.shape(0)),
+      /* const int  C = */ static_cast<int>(in_padded.shape(3)),
+      /* const int  O = */ static_cast<int>(wt.shape(0)),
+      /* const int iS[NDIM] = */
+      {static_cast<int>(in_padded.shape(1)),
+       static_cast<int>(in_padded.shape(2))},
+      /* const int wS[NDIM] = */
+      {static_cast<int>(wt.shape(1)), static_cast<int>(wt.shape(2))},
+      /* const int oS[NDIM] = */
+      {static_cast<int>(out.shape(1)), static_cast<int>(out.shape(2))},
+      /* const int str[NDIM] = */ {1, 1},
+      /* const int pad[NDIM] = */ {0, 0},
+      /* const int kdil[NDIM] = */ {1, 1},
+      /* const int idil[NDIM] = */ {1, 1},
+      /* const size_t in_strides[NDIM + 2] = */
+      {in_padded.strides()[0],
+       in_padded.strides()[1],
+       in_padded.strides()[2],
+       in_padded.strides()[3]},
+      /* const size_t wt_strides[NDIM + 2] = */
+      {wt.strides()[0], wt.strides()[1], wt.strides()[2], wt.strides()[3]},
+      /* const size_t out_strides[NDIM + 2] = */
+      {out.strides()[0], out.strides()[1], out.strides()[2], out.strides()[3]},
+      /* const int groups = */ 1,
+      /* const bool flip = */ false,
+  };
+
   int O_c = conv_params.O;
   int C_c = conv_params.C;
 
@@ -598,7 +620,7 @@ void winograd_conv_2D_gpu(
     int bo = 4;
     std::ostringstream kname;
     kname << "winograd_conv_2d_weight_transform_" << type_to_name(out) << "_bc"
-          << bc << "_flip" << conv_params.flip;
+          << bc;
     auto& compute_encoder = d.get_command_encoder(s.index);
     auto kernel = d.get_kernel(kname.str());
     compute_encoder.set_compute_pipeline_state(kernel);
@@ -631,10 +653,10 @@ void winograd_conv_2D_gpu(
     auto kernel = d.get_kernel(kname.str());
     compute_encoder.set_compute_pipeline_state(kernel);
 
-    compute_encoder.set_input_array(in, 0);
+    compute_encoder.set_input_array(in_padded, 0);
     compute_encoder.set_output_array(inp_wg, 1);
 
-    compute_encoder.set_bytes(conv_params, 2);
+    compute_encoder.set_bytes(conv_params_updated, 2);
 
     MTL::Size group_dims = MTL::Size(32, wn, wm);
     MTL::Size grid_dims = MTL::Size(N_tiles_w, N_tiles_h, N_tiles_n);
@@ -681,7 +703,7 @@ void winograd_conv_2D_gpu(
     compute_encoder.set_input_array(out_wg, 0);
     compute_encoder.set_output_array(out, 1);
 
-    compute_encoder.set_bytes(conv_params, 2);
+    compute_encoder.set_bytes(conv_params_updated, 2);
 
     MTL::Size group_dims = MTL::Size(32, wn, wm);
     MTL::Size grid_dims = MTL::Size(N_tiles_w, N_tiles_h, N_tiles_n);
@@ -745,18 +767,14 @@ void conv_2D_gpu(
   }
 
   // Direct to winograd conv
-  bool img_large =
+  bool inp_large =
       (conv_params.N * conv_params.iS[0] * conv_params.iS[1]) >= 1ul << 12;
   bool channels_large = (conv_params.C + conv_params.O) >= 256;
-  if (conv_params.wS[0] == 3 && conv_params.wS[1] == 3 &&
-      conv_params.C % 32 == 0 && conv_params.O % 32 == 0 && is_stride_one &&
-      is_kdil_one && is_idil_one) {
-    if (img_large && channels_large) {
-      return winograd_conv_2D_gpu(s, d, in, wt, out, conv_params, copies);
-    }
-    if (conv_params.N <= 1) {
-      return winograd_conv_2D_fused_gpu(s, d, in, wt, out, conv_params, copies);
-    }
+  if (!flip && is_stride_one && is_kdil_one && is_idil_one &&
+      conv_params.wS[0] == 3 && conv_params.wS[1] == 3 &&
+      conv_params.C % 32 == 0 && conv_params.O % 32 == 0 && inp_large &&
+      channels_large) {
+    return winograd_conv_2D_gpu(s, d, in, wt, out, conv_params, copies);
   }
 
   // Direct to implicit gemm conv
@@ -858,40 +876,8 @@ void Convolution::eval_gpu(const std::vector<array>& inputs, array& out) {
     wt = arr_copy;
   }
 
-  // Check for 1x1 conv
-  auto is_one = [](int x) { return x == 1; };
-  auto is_zero = [](int x) { return x == 0; };
-  if (groups_ == 1 && (wt.shape(0) * wt.shape(-1) == wt.size()) &&
-      std::all_of(wt.shape().begin() + 1, wt.shape().end() - 1, is_one) &&
-      std::all_of(kernel_strides_.begin(), kernel_strides_.end(), is_one) &&
-      std::all_of(input_dilation_.begin(), input_dilation_.end(), is_one) &&
-      std::all_of(kernel_dilation_.begin(), kernel_dilation_.end(), is_one) &&
-      std::all_of(padding_.begin(), padding_.end(), is_zero)) {
-    std::vector<array> empty_copies;
-    steel_matmul_regular(
-        s,
-        d,
-        /*a = */ in,
-        /*b = */ wt,
-        /*c = */ out,
-        /*M = */ in.size() / in.shape(-1),
-        /*N = */ wt.shape(0),
-        /*K = */ in.shape(-1),
-        /*batch_size_out = */ 1,
-        /*lda = */ in.shape(-1),
-        /*ldb = */ wt.shape(-1),
-        /*ldd = */ wt.shape(0),
-        /*transpose_a = */ false,
-        /*transpose_b = */ true,
-        /*batch_shape = */ {1},
-        /*batch_strides = */ {1},
-        /*A_batch_stride = */ 0,
-        /*B_batch_stride = */ 0,
-        /*matrix_stride_out = */ 0,
-        /*copies = */ empty_copies);
-  }
   // 3D conv
-  else if (out.ndim() == 5) {
+  if (out.ndim() == 5) {
     conv_3D_gpu(
         s,
         d,

mlx/backend/metal/device.cpp

@@ -249,14 +249,12 @@ Device::~Device() {
   for (auto& l : library_map_) {
     l.second->release();
   }
+  stream_map_.clear();
   device_->release();
 }
 
 void Device::new_queue(int index) {
   auto thread_pool = metal::new_scoped_memory_pool();
-
-  // Multiple threads can ask the device for queues
-  // We lock this as a critical section for safety
   auto q = device_->newCommandQueue(MAX_BUFFERS_PER_QUEUE);
   debug_set_stream_queue_label(q, index);
   if (!q) {
@@ -269,6 +267,10 @@ void Device::new_queue(int index) {
   }
 }
 
+MTL::CommandQueue* Device::get_queue(Stream stream) {
+  return get_stream_(stream.index).queue;
+}
+
 bool Device::command_buffer_needs_commit(int index) {
   auto& stream = get_stream_(index);
   if (stream.buffer_ops > max_ops_per_buffer_ ||
@@ -690,12 +692,13 @@ void new_stream(Stream stream) {
   }
 }
 
-std::unordered_map<std::string, std::variant<std::string, size_t>>
+const std::unordered_map<std::string, std::variant<std::string, size_t>>&
 device_info() {
   auto init_device_info = []()
       -> std::unordered_map<std::string, std::variant<std::string, size_t>> {
     auto pool = new_scoped_memory_pool();
     auto raw_device = device(default_device()).mtl_device();
+    auto name = std::string(raw_device->name()->utf8String());
     auto arch = std::string(raw_device->architecture()->name()->utf8String());
 
     size_t memsize = 0;
@@ -709,6 +712,7 @@ device_info() {
     }
 
     return {
+        {"device_name", name},
         {"architecture", arch},
         {"max_buffer_length", raw_device->maxBufferLength()},
         {"max_recommended_working_set_size",

mlx/backend/metal/device.h

@@ -178,6 +178,9 @@ class Device {
   }
 
   void new_queue(int index);
+
+  MTL::CommandQueue* get_queue(Stream stream);
+
   MTL::CommandBuffer* get_command_buffer(int index);
   bool command_buffer_needs_commit(int index);
   void commit_command_buffer(int index);

mlx/backend/metal/kernels/CMakeLists.txt

@@ -4,11 +4,12 @@ set(BASE_HEADERS
     bf16_math.h
     complex.h
     defines.h
+    erf.h
     expm1f.h
     utils.h)
 
 function(build_kernel_base TARGET SRCFILE DEPS)
-  set(METAL_FLAGS -Wall -Wextra -fno-fast-math)
+  set(METAL_FLAGS -Wall -Wextra -fno-fast-math -Wno-c++17-extensions)
   if(MLX_METAL_DEBUG)
     set(METAL_FLAGS ${METAL_FLAGS} -gline-tables-only -frecord-sources)
   endif()

mlx/backend/metal/kernels/conv.metal

@@ -326,13 +326,7 @@ constant constexpr const float WinogradTransforms<6, 3, 8>::wt_transform[8][8];
 constant constexpr const float WinogradTransforms<6, 3, 8>::in_transform[8][8];
 constant constexpr const float WinogradTransforms<6, 3, 8>::out_transform[8][8];
 
-template <
-    typename T,
-    int BC = 32,
-    int BO = 4,
-    bool do_flip = false,
-    int M = 6,
-    int R = 3>
+template <typename T, int BC = 32, int BO = 4, int M = 6, int R = 3>
 [[kernel, max_total_threads_per_threadgroup(BO * 32)]] void
 winograd_conv_2d_weight_transform(
     const device T* wt_in [[buffer(0)]],
@@ -379,12 +373,7 @@ winograd_conv_2d_weight_transform(
     for (int kh = 0; kh < R; ++kh) {
       for (int kw = 0; kw < R; ++kw) {
         for (int kc = simd_lane_id; kc < BC; kc += 32) {
-          if (do_flip) {
-            Ws[simd_group_id][R - 1 - kh][R - 1 - kw][kc] =
-                wt_in[kh * R * C + kw * C + kc];
-          } else {
-            Ws[simd_group_id][kh][kw][kc] = wt_in[kh * R * C + kw * C + kc];
-          }
+          Ws[simd_group_id][kh][kw][kc] = wt_in[kh * R * C + kw * C + kc];
         }
       }
     }
@@ -409,10 +398,10 @@ winograd_conv_2d_weight_transform(
   }
 }
 
-#define instantiate_winograd_conv_2d_weight_tr_base_2(name, itype, bc, f)   \
-  template [[host_name("winograd_conv_2d_weight_transform_" #name "_bc" #bc \
-                       "_flip" #f)]] [[kernel]] void                        \
-  winograd_conv_2d_weight_transform<itype, bc, 4, f>(                       \
+#define instantiate_winograd_conv_2d_weight_transform_base(name, itype, bc) \
+  template [[host_name("winograd_conv_2d_weight_transform_" #name           \
+                       "_bc" #bc)]] [[kernel]] void                         \
+  winograd_conv_2d_weight_transform<itype, bc>(                             \
       const device itype* wt_in [[buffer(0)]],                              \
       device itype* wt_out [[buffer(1)]],                                   \
       const constant int& C [[buffer(2)]],                                  \
@@ -421,10 +410,6 @@ winograd_conv_2d_weight_transform(
       uint simd_group_id [[simdgroup_index_in_threadgroup]],                \
       uint simd_lane_id [[thread_index_in_simdgroup]]);
 
-#define instantiate_winograd_conv_2d_weight_transform_base(name, itype, bc) \
-  instantiate_winograd_conv_2d_weight_tr_base_2(name, itype, bc, 0)         \
-      instantiate_winograd_conv_2d_weight_tr_base_2(name, itype, bc, 1)
-
 template <typename T, int BC, int WM, int WN, int M = 6, int R = 3>
 [[kernel, max_total_threads_per_threadgroup(WM* WN * 32)]] void
 winograd_conv_2d_input_transform(
@@ -460,17 +445,10 @@ winograd_conv_2d_input_transform(
   // Resolve input tile
   constexpr int TH = (A / WM);
   constexpr int TW = (A / WN);
-  const int kh = TH * (simd_group_id / WN);
-  const int kw = TW * (simd_group_id % WN);
-  const int bh = M * tid.y + kh - params.pad[1];
-  const int bw = M * tid.x + kw - params.pad[0];
-
-  const bool is_edge_w_lo = bw < 0;
-  const bool is_edge_h_lo = bh < 0;
-  const bool is_edge_w_hi = bw + (TW - 1) >= params.iS[0];
-  const bool is_edge_h_hi = bh + (TH - 1) >= params.iS[1];
-  const bool is_edge =
-      is_edge_w_lo || is_edge_h_lo || is_edge_w_hi || is_edge_h_hi;
+  int kh = TH * (simd_group_id / WN);
+  int kw = TW * (simd_group_id % WN);
+  int bh = M * tid.y + kh;
+  int bw = M * tid.x + kw;
 
   // Move to the correct input tile
   inp_in += tid.z * params.in_strides[0] + bh * params.in_strides[1] +
@@ -506,21 +484,8 @@ winograd_conv_2d_input_transform(
     for (int h = 0; h < TH; h++) {
       for (int w = 0; w < TW; w++) {
         const device T* in_ptr = inp_in + jump_in[h][w];
-        if (is_edge) {
-          if (((bh + h) < 0 || (bh + h) >= params.iS[1]) ||
-              ((bw + w) < 0 || (bw + w) >= params.iS[0])) {
-            for (int c = simd_lane_id; c < BC; c += 32) {
-              Is[kh + h][kw + w][c] = T(0);
-            }
-          } else {
-            for (int c = simd_lane_id; c < BC; c += 32) {
-              Is[kh + h][kw + w][c] = in_ptr[c];
-            }
-          }
-        } else {
-          for (int c = simd_lane_id; c < BC; c += 32) {
-            Is[kh + h][kw + w][c] = in_ptr[c];
-          }
+        for (int c = simd_lane_id; c < BC; c += 32) {
+          Is[kh + h][kw + w][c] = in_ptr[c];
         }
       }
     }
@@ -687,373 +652,3 @@ winograd_conv_2d_output_transform(
 instantiate_winograd_conv_2d(float32, float);
 instantiate_winograd_conv_2d(bfloat16, bfloat16_t);
 instantiate_winograd_conv_2d(float16, half); // clang-format on
-
-#include "mlx/backend/metal/kernels/steel/attn/mma.h"
-
-template <
-    typename T,
-    bool do_flip = false,
-    int WM = 4,
-    int WN = 1,
-    typename AccumType = float>
-[[kernel]] void winograd_fused(
-    const device T* input [[buffer(0)]],
-    const device T* weight [[buffer(1)]],
-    device T* output [[buffer(2)]],
-    const constant MLXConvParams<2>& params [[buffer(3)]],
-    uint3 tid [[threadgroup_position_in_grid]],
-    uint3 lid [[thread_position_in_threadgroup]],
-    uint3 tgp_per_grid [[threadgroups_per_grid]],
-    ushort simd_group_id [[simdgroup_index_in_threadgroup]],
-    ushort simd_lane_id [[thread_index_in_simdgroup]]) {
-  using namespace mlx::steel;
-
-  (void)tgp_per_grid;
-
-  // Winograd F(n x n, r x r)
-  // n x n output window
-  constexpr short FN = 2;
-  // r x r filter size
-  constexpr short FR = 3;
-  // a x a input window, a = n + r - 1
-  constexpr short FA = 4;
-
-  constexpr short kFragSize = 8; // MMA frag size
-
-  constexpr short BT = 8; // Tile block size
-  constexpr short BO = 8; // Output channel block size
-  constexpr short BC = 8; // Input channel block size
-
-  // clang-format off
-  static_assert(BT % (1 * kFragSize) == 0 && 
-                BO % (1 * kFragSize) == 0 && 
-                BC % kFragSize == 0,
-                "Matmuls sizes must be compatible with fragments");
-  // clang-format on
-
-  // Prepare for matmul
-
-  // Warp tile sizes for matmul
-  constexpr short TM = (FA * FA * BT) / (WM * kFragSize);
-  constexpr short TN = (BO) / (WN * kFragSize);
-  constexpr short TK = (BC) / (kFragSize);
-
-  // Warp primitives
-  using MMAFrag_acc_t = BaseMMAFrag<AccumType, kFragSize, kFragSize>;
-
-  // Warp tiles sizes for matmul
-  MMATile<AccumType, 1, TK, MMAFrag_acc_t> Itile;
-  MMATile<AccumType, TK, TN, MMAFrag_acc_t> Wtile;
-  MMATile<AccumType, 1, TN, MMAFrag_acc_t> Otile[TM];
-
-  for (int im = 0; im < 4; im++) {
-    Otile[im].clear();
-  }
-
-  // Threadgroup memory for Weights and Inputs
-  constexpr short BS = BT > BO ? BT : BO;
-  threadgroup T Wt[FA * FA * BC * BO];
-  threadgroup T It[FA * FA * BS * BS];
-
-  // Get thread position in tile
-  short2 simd_coord = MMAFrag_acc_t::get_coord(simd_lane_id);
-  const short sm = simd_coord.y;
-  const short sn = simd_coord.x;
-
-  static_assert(FA * FA * BT == 32 * WM * WN, "Each thread loads one pixel.");
-  const int thr_idx = simd_group_id * 32 + simd_lane_id;
-  const int thr_t = thr_idx / (FA * FA);
-  const int thr_hw = thr_idx % (FA * FA);
-  const int thr_h = thr_hw / FA;
-  const int thr_w = thr_hw % FA;
-
-  // Get batch, tile, and output idx for warp
-  const int b_idx = tid.z;
-  const int t_idx = BT * tid.y + thr_t;
-  const int o_idx = BO * tid.x + thr_t;
-
-  // Divide tile into h, w tile
-  uniform<int> oHu = make_uniform(params.oS[0]);
-  uniform<int> oWu = make_uniform(params.oS[1]);
-  uniform<int> tHu = (oHu + make_uniform(FN - 1)) / make_uniform(FN);
-  uniform<int> tWu = (oWu + make_uniform(FN - 1)) / make_uniform(FN);
-
-  const int oH_idx = FN * (t_idx / tWu);
-  const int oW_idx = FN * (t_idx % tWu);
-  const int iH_idx = oH_idx + thr_h - params.pad[0];
-  const int iW_idx = oW_idx + thr_w - params.pad[1];
-
-  // Move to correct location
-
-  // clang-format off
-  input +=  b_idx * params.in_strides[0] + // N
-           iH_idx * params.in_strides[1] + // H
-           iW_idx * params.in_strides[2];  // W
-
-  weight += o_idx * params.wt_strides[0] + // O
-            thr_h * params.wt_strides[1] + // H 
-            thr_w * params.wt_strides[2];  // W
-  // clang-format on
-
-  // Do edge check prep for input
-  const bool is_edge_w_lo = iH_idx < 0;
-  const bool is_edge_h_lo = iW_idx < 0;
-  const bool is_edge_w_hi = iH_idx >= params.iS[0];
-  const bool is_edge_h_hi = iW_idx >= params.iS[1];
-  const bool is_edge =
-      is_edge_w_lo || is_edge_h_lo || is_edge_w_hi || is_edge_h_hi;
-
-  // Iterate over C
-  for (int c = 0; c < params.C; c += BC) {
-#define tmp_load_wt_idx(o, h, w, c) h* FA* BC* BO + w* BC* BO + c* BO + o
-#define tmp_load_in_idx(t, h, w, c) h* FA* BS* BC + w* BS* BC + t* BC + c
-
-#define tmp_trns_wt_idx(o, h, w, c) h* FA* BC* BO + w* BC* BO + c* BO + o
-#define tmp_trns_in_idx(t, h, w, c) h* FA* BS* BC + w* BS* BC + t* BC + c
-
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-
-    // Load weight
-    if (thr_h < FR && thr_w < FR && thr_t < BO) {
-      for (int ic = 0; ic < BC; ic++) {
-        if (do_flip) {
-          Wt[tmp_load_wt_idx(thr_t, FR - 1 - thr_h, FR - 1 - thr_w, ic)] =
-              weight[c + ic];
-        } else {
-          Wt[tmp_load_wt_idx(thr_t, thr_h, thr_w, ic)] = weight[c + ic];
-        }
-      }
-    }
-
-    // Load input
-    if (is_edge) {
-      for (int ic = 0; ic < BC; ic++) {
-        It[tmp_load_in_idx(thr_t, thr_h, thr_w, ic)] = T(0);
-      }
-    } else {
-      for (int ic = 0; ic < BC; ic++) {
-        It[tmp_load_in_idx(thr_t, thr_h, thr_w, ic)] = input[c + ic];
-      }
-    }
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-
-    // Transform weight
-    if (lid.z == 0) {
-      const short ic = lid.y;
-      const short io = lid.x;
-
-      T tmp_0[4][4];
-      T tmp_1[4][4];
-
-      for (int ii = 0; ii < 3; ++ii) {
-        for (int jj = 0; jj < 3; ++jj) {
-          tmp_0[ii][jj] = Wt[tmp_load_wt_idx(io, ii, jj, ic)];
-        }
-      }
-
-      //////////////////////////////////////////////
-
-      tmp_1[0][0] = tmp_0[0][0];
-      tmp_1[0][1] = tmp_0[0][1];
-      tmp_1[0][2] = tmp_0[0][2];
-
-      tmp_1[1][0] = T(0.5) * (tmp_0[0][0] + tmp_0[1][0] + tmp_0[2][0]);
-      tmp_1[1][1] = T(0.5) * (tmp_0[0][1] + tmp_0[1][1] + tmp_0[2][1]);
-      tmp_1[1][2] = T(0.5) * (tmp_0[0][2] + tmp_0[1][2] + tmp_0[2][2]);
-
-      tmp_1[2][0] = tmp_1[1][0] - tmp_0[1][0];
-      tmp_1[2][1] = tmp_1[1][1] - tmp_0[1][1];
-      tmp_1[2][2] = tmp_1[1][2] - tmp_0[1][2];
-
-      tmp_1[3][0] = tmp_0[2][0];
-      tmp_1[3][1] = tmp_0[2][1];
-      tmp_1[3][2] = tmp_0[2][2];
-
-      //////////////////////////////////////////////
-      tmp_0[0][0] = tmp_1[0][0];
-      tmp_0[1][0] = tmp_1[1][0];
-      tmp_0[2][0] = tmp_1[2][0];
-      tmp_0[3][0] = tmp_1[3][0];
-
-      tmp_0[0][1] = T(0.5) * (tmp_1[0][0] + tmp_1[0][1] + tmp_1[0][2]);
-      tmp_0[1][1] = T(0.5) * (tmp_1[1][0] + tmp_1[1][1] + tmp_1[1][2]);
-      tmp_0[2][1] = T(0.5) * (tmp_1[2][0] + tmp_1[2][1] + tmp_1[2][2]);
-      tmp_0[3][1] = T(0.5) * (tmp_1[3][0] + tmp_1[3][1] + tmp_1[3][2]);
-
-      tmp_0[0][2] = tmp_0[0][1] - tmp_1[0][1];
-      tmp_0[1][2] = tmp_0[1][1] - tmp_1[1][1];
-      tmp_0[2][2] = tmp_0[2][1] - tmp_1[2][1];
-      tmp_0[3][2] = tmp_0[3][1] - tmp_1[3][1];
-
-      tmp_0[0][3] = tmp_1[0][2];
-      tmp_0[1][3] = tmp_1[1][2];
-      tmp_0[2][3] = tmp_1[2][2];
-      tmp_0[3][3] = tmp_1[3][2];
-
-      for (int ii = 0; ii < 4; ++ii) {
-        for (int jj = 0; jj < 4; ++jj) {
-          Wt[tmp_trns_wt_idx(io, ii, jj, ic)] = tmp_0[ii][jj];
-        }
-      }
-    }
-
-    // Transform input
-    else {
-      const short it = lid.y;
-      const short ic = lid.x;
-
-      T tmp_0[4][4];
-      T tmp_1[4][4];
-
-      for (int ii = 0; ii < 4; ++ii) {
-        for (int jj = 0; jj < 4; ++jj) {
-          tmp_0[ii][jj] = It[tmp_load_in_idx(it, ii, jj, ic)];
-        }
-      }
-
-      //////////////////////////////////////////////
-
-      tmp_1[0][0] = tmp_0[0][0] - tmp_0[2][0];
-      tmp_1[0][1] = tmp_0[0][1] - tmp_0[2][1];
-      tmp_1[0][2] = tmp_0[0][2] - tmp_0[2][2];
-      tmp_1[0][3] = tmp_0[0][3] - tmp_0[2][3];
-
-      tmp_1[1][0] = tmp_0[1][0] + tmp_0[2][0];
-      tmp_1[1][1] = tmp_0[1][1] + tmp_0[2][1];
-      tmp_1[1][2] = tmp_0[1][2] + tmp_0[2][2];
-      tmp_1[1][3] = tmp_0[1][3] + tmp_0[2][3];
-
-      tmp_1[2][0] = tmp_0[2][0] - tmp_0[1][0];
-      tmp_1[2][1] = tmp_0[2][1] - tmp_0[1][1];
-      tmp_1[2][2] = tmp_0[2][2] - tmp_0[1][2];
-      tmp_1[2][3] = tmp_0[2][3] - tmp_0[1][3];
-
-      tmp_1[3][0] = tmp_0[1][0] - tmp_0[3][0];
-      tmp_1[3][1] = tmp_0[1][1] - tmp_0[3][1];
-      tmp_1[3][2] = tmp_0[1][2] - tmp_0[3][2];
-      tmp_1[3][3] = tmp_0[1][3] - tmp_0[3][3];
-
-      //////////////////////////////////////////////
-      tmp_0[0][0] = tmp_1[0][0] - tmp_1[0][2];
-      tmp_0[1][0] = tmp_1[1][0] - tmp_1[1][2];
-      tmp_0[2][0] = tmp_1[2][0] - tmp_1[2][2];
-      tmp_0[3][0] = tmp_1[3][0] - tmp_1[3][2];
-
-      tmp_0[0][1] = tmp_1[0][1] + tmp_1[0][2];
-      tmp_0[1][1] = tmp_1[1][1] + tmp_1[1][2];
-      tmp_0[2][1] = tmp_1[2][1] + tmp_1[2][2];
-      tmp_0[3][1] = tmp_1[3][1] + tmp_1[3][2];
-
-      tmp_0[0][2] = tmp_1[0][2] - tmp_1[0][1];
-      tmp_0[1][2] = tmp_1[1][2] - tmp_1[1][1];
-      tmp_0[2][2] = tmp_1[2][2] - tmp_1[2][1];
-      tmp_0[3][2] = tmp_1[3][2] - tmp_1[3][1];
-
-      tmp_0[0][3] = tmp_1[0][1] - tmp_1[0][3];
-      tmp_0[1][3] = tmp_1[1][1] - tmp_1[1][3];
-      tmp_0[2][3] = tmp_1[2][1] - tmp_1[2][3];
-      tmp_0[3][3] = tmp_1[3][1] - tmp_1[3][3];
-
-      for (int ii = 0; ii < 4; ++ii) {
-        for (int jj = 0; jj < 4; ++jj) {
-          It[tmp_trns_in_idx(it, ii, jj, ic)] = tmp_0[ii][jj];
-        }
-      }
-    }
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-
-    // Do matmul
-    for (int im = 0; im < 4; im++) {
-      simdgroup_barrier(mem_flags::mem_none);
-      Itile.template load<T, 1, 1, BS, 1>(
-          &It[simd_group_id * FA * BS * BS + im * BS * BS + sm * BS + sn]);
-      simdgroup_barrier(mem_flags::mem_none);
-      Wtile.template load<T, 1, 1, BO, 1>(
-          &Wt[simd_group_id * FA * BC * BO + im * BC * BO + sm * BO + sn]);
-      simdgroup_barrier(mem_flags::mem_none);
-      tile_matmad(Otile[im], Itile, Wtile, Otile[im]);
-    }
-  }
-
-  // Transform and write output
-  threadgroup_barrier(mem_flags::mem_threadgroup);
-  for (int im = 0; im < 4; im++) {
-    Otile[im].template store<T, 1, 1, BS, 1>(
-        &It[simd_group_id * FA * BS * BS + im * BS * BS + sm * BS + sn]);
-  }
-  threadgroup_barrier(mem_flags::mem_threadgroup);
-
-  if (lid.z == 0) {
-    const short it = lid.y;
-    const short io = lid.x;
-
-    T tmp_0[4][4];
-    T tmp_1[2][4];
-    T tmp_2[2][2];
-
-    for (int ii = 0; ii < 4; ++ii) {
-      for (int jj = 0; jj < 4; ++jj) {
-        tmp_0[ii][jj] = It[tmp_trns_in_idx(it, ii, jj, io)];
-      }
-    }
-
-    tmp_1[0][0] = tmp_0[0][0] + tmp_0[1][0] + tmp_0[2][0];
-    tmp_1[0][1] = tmp_0[0][1] + tmp_0[1][1] + tmp_0[2][1];
-    tmp_1[0][2] = tmp_0[0][2] + tmp_0[1][2] + tmp_0[2][2];
-    tmp_1[0][3] = tmp_0[0][3] + tmp_0[1][3] + tmp_0[2][3];
-
-    tmp_1[1][0] = tmp_0[1][0] - tmp_0[2][0] - tmp_0[3][0];
-    tmp_1[1][1] = tmp_0[1][1] - tmp_0[2][1] - tmp_0[3][1];
-    tmp_1[1][2] = tmp_0[1][2] - tmp_0[2][2] - tmp_0[3][2];
-    tmp_1[1][3] = tmp_0[1][3] - tmp_0[2][3] - tmp_0[3][3];
-
-    tmp_2[0][0] = tmp_1[0][0] + tmp_1[0][1] + tmp_1[0][2];
-    tmp_2[1][0] = tmp_1[1][0] + tmp_1[1][1] + tmp_1[1][2];
-
-    tmp_2[0][1] = tmp_1[0][1] - tmp_1[0][2] - tmp_1[0][3];
-    tmp_2[1][1] = tmp_1[1][1] - tmp_1[1][2] - tmp_1[1][3];
-
-    const int oH_i = FN * ((BT * tid.y + it) / tWu);
-    const int oW_i = FN * ((BT * tid.y + it) % tWu);
-
-    // clang-format off
-    output += b_idx * params.out_strides[0] + // N
-              oH_i * params.out_strides[1] +  // H
-              oW_i * params.out_strides[2] +  // W
-              BO * tid.x;                     // C
-
-    // clang-format on
-
-    output[0 * params.out_strides[1] + 0 * params.out_strides[2] + io] =
-        tmp_2[0][0];
-    output[0 * params.out_strides[1] + 1 * params.out_strides[2] + io] =
-        tmp_2[0][1];
-    output[1 * params.out_strides[1] + 0 * params.out_strides[2] + io] =
-        tmp_2[1][0];
-    output[1 * params.out_strides[1] + 1 * params.out_strides[2] + io] =
-        tmp_2[1][1];
-  }
-}
-
-// clang-format off
-#define instantiate_winograd_conv_2d_fused(name, itype, f)            \
-  template [[host_name("winograd_conv_2d_fused_" #name "_flip" #f)]]  \
-  [[kernel]] void winograd_fused<itype, f>(                           \
-      const device itype* input [[buffer(0)]],                        \
-      const device itype* weight [[buffer(1)]],                       \
-      device itype* output [[buffer(2)]],                             \
-      const constant MLXConvParams<2>& params [[buffer(3)]],          \
-      uint3 tid [[threadgroup_position_in_grid]],                     \
-      uint3 lid [[thread_position_in_threadgroup]],                   \
-      uint3 tgp_per_grid [[threadgroups_per_grid]],                   \
-      ushort simd_group_id [[simdgroup_index_in_threadgroup]],        \
-      ushort simd_lane_id [[thread_index_in_simdgroup]]);
-
-#define instantiate_winograd_conv_2d_fused_2(name, itype) \
-  instantiate_winograd_conv_2d_fused(name, itype, 0)      \
-  instantiate_winograd_conv_2d_fused(name, itype, 1)
-
-instantiate_winograd_conv_2d_fused_2(float32, float);
-instantiate_winograd_conv_2d_fused_2(float16, float16_t);
-instantiate_winograd_conv_2d_fused_2(bfloat16, bfloat16_t);
-// clang-format on

mlx/backend/metal/kernels/layer_norm.metal

@@ -7,6 +7,8 @@
 
 using namespace metal;
 
+constant bool has_w [[function_constant(20)]];
+
 template <typename T, int N_READS = RMS_N_READS>
 [[kernel]] void layer_norm_single_row(
     const device T* x,
@@ -327,7 +329,9 @@ template <typename T, int N_READS = RMS_N_READS>
       gx[i] = static_cast<T>(
           normalizer * (thread_w[i] * thread_g[i] - meanwg) -
           thread_x[i] * meanwgxc * normalizer2);
-      gw[i] = static_cast<T>(thread_g[i] * thread_x[i]);
+      if (has_w) {
+        gw[i] = static_cast<T>(thread_g[i] * thread_x[i]);
+      }
     }
   } else {
     for (int i = 0; i < N_READS; i++) {
@@ -336,7 +340,9 @@ template <typename T, int N_READS = RMS_N_READS>
         gx[i] = static_cast<T>(
             normalizer * (thread_w[i] * thread_g[i] - meanwg) -
             thread_x[i] * meanwgxc * normalizer2);
-        gw[i] = static_cast<T>(thread_g[i] * thread_x[i]);
+        if (has_w) {
+          gw[i] = static_cast<T>(thread_g[i] * thread_x[i]);
+        }
       }
     }
   }
@@ -465,7 +471,9 @@ template <typename T, int N_READS = RMS_N_READS>
         float gi = g[i + r];
         gx[i + r] = static_cast<T>(
             normalizer * (wi * gi - meanwg) - xi * meanwgxc * normalizer2);
-        gw[i + r] = static_cast<T>(gi * xi);
+        if (has_w) {
+          gw[i + r] = static_cast<T>(gi * xi);
+        }
       }
     } else {
       for (int i = 0; i < N_READS; i++) {
@@ -475,7 +483,9 @@ template <typename T, int N_READS = RMS_N_READS>
           float gi = g[i + r];
           gx[i + r] = static_cast<T>(
               normalizer * (wi * gi - meanwg) - xi * meanwgxc * normalizer2);
-          gw[i + r] = static_cast<T>(gi * xi);
+          if (has_w) {
+            gw[i + r] = static_cast<T>(gi * xi);
+          }
         }
       }
     }

mlx/backend/metal/kernels/quantized.h

@@ -2015,9 +2015,9 @@ template <typename T, const int group_size, const int bits>
     device T* biases [[buffer(3)]],
     uint2 index [[thread_position_in_grid]],
     uint2 grid_dim [[threads_per_grid]]) {
-  constexpr T eps = T(1e-7);
+  constexpr float eps = 1e-7;
   constexpr int simd_size = 32;
-  constexpr T n_bins = (1 << bits) - 1;
+  constexpr float n_bins = (1 << bits) - 1;
   constexpr int packs_per_int = bits == 3 ? 8 : bits == 6 ? 4 : 8 / bits;
   constexpr int values_per_reduce = group_size / simd_size;
   constexpr int writes_per_reduce = packs_per_int / values_per_reduce;
@@ -2036,13 +2036,13 @@ template <typename T, const int group_size, const int bits>
       ? offset * writes_per_pack
       : offset * bytes_per_pack / writes_per_reduce;
 
-  T w_thread[values_per_reduce];
-  T w_min = Limits<T>::max;
-  T w_max = 0;
+  float w_thread[values_per_reduce];
+  float w_min = Limits<T>::max;
+  float w_max = 0;
 
 #pragma clang loop unroll(full)
   for (int i = 0; i < values_per_reduce; i++) {
-    T val = w[in_index + i];
+    float val = w[in_index + i];
     w_thread[i] = val;
     w_min = min(w_min, val);
     w_max = max(w_max, val);
@@ -2051,20 +2051,20 @@ template <typename T, const int group_size, const int bits>
   w_min = simd_min(w_min);
   w_max = simd_max(w_max);
 
-  T scale = max((w_max - w_min) / n_bins, eps);
+  float scale = max((w_max - w_min) / n_bins, eps);
   bool side = abs(w_min) > abs(w_max);
   scale = side ? scale : -scale;
-  T edge = side ? w_min : w_max;
-  T q0 = round(edge / scale);
+  float edge = side ? w_min : w_max;
+  float q0 = round(edge / scale);
   bool at_zero = q0 == 0.0f;
   scale = at_zero ? scale : edge / q0;
-  T bias = at_zero ? T(0) : edge;
+  float bias = at_zero ? 0 : edge;
 
   // Write out the scales and biases
   size_t gindex = in_index / group_size;
   if (in_index % group_size == 0) {
-    scales[gindex] = scale;
-    biases[gindex] = bias;
+    scales[gindex] = static_cast<T>(scale);
+    biases[gindex] = static_cast<T>(bias);
   }
 
   // We accumulate 3 bytes worth for 3/6 bit so we need a uint32_t

mlx/backend/metal/kernels/rms_norm.metal

@@ -7,6 +7,8 @@
 
 using namespace metal;
 
+constant bool has_w [[function_constant(20)]];
+
 template <typename T, int N_READS = RMS_N_READS>
 [[kernel]] void rms_single_row(
     const device T* x,
@@ -243,7 +245,9 @@ template <typename T, int N_READS = RMS_N_READS>
       gx[i] = static_cast<T>(
           thread_g[i] * thread_w[i] * normalizer -
           thread_x[i] * meangwx * normalizer3);
-      gw[i] = static_cast<T>(thread_g[i] * thread_x[i] * normalizer);
+      if (has_w) {
+        gw[i] = static_cast<T>(thread_g[i] * thread_x[i] * normalizer);
+      }
     }
   } else {
     for (int i = 0; i < N_READS; i++) {
@@ -251,7 +255,9 @@ template <typename T, int N_READS = RMS_N_READS>
         gx[i] = static_cast<T>(
             thread_g[i] * thread_w[i] * normalizer -
             thread_x[i] * meangwx * normalizer3);
-        gw[i] = static_cast<T>(thread_g[i] * thread_x[i] * normalizer);
+        if (has_w) {
+          gw[i] = static_cast<T>(thread_g[i] * thread_x[i] * normalizer);
+        }
       }
     }
   }
@@ -351,7 +357,9 @@ template <typename T, int N_READS = RMS_N_READS>
 
         gx[i + r] =
             static_cast<T>(gi * wi * normalizer - xi * meangwx * normalizer3);
-        gw[i + r] = static_cast<T>(gi * xi * normalizer);
+        if (has_w) {
+          gw[i + r] = static_cast<T>(gi * xi * normalizer);
+        }
       }
     } else {
       for (int i = 0; i < N_READS; i++) {
@@ -362,7 +370,9 @@ template <typename T, int N_READS = RMS_N_READS>
 
           gx[i + r] =
               static_cast<T>(gi * wi * normalizer - xi * meangwx * normalizer3);
-          gw[i + r] = static_cast<T>(gi * xi * normalizer);
+          if (has_w) {
+            gw[i + r] = static_cast<T>(gi * xi * normalizer);
+          }
         }
       }
     }

mlx/backend/metal/kernels/sdpa_vector.h

@@ -5,6 +5,7 @@
 using namespace metal;
 
 constant bool has_mask [[function_constant(20)]];
+constant bool query_transposed [[function_constant(21)]];
 
 template <typename T, int D, int V = D>
 [[kernel]] void sdpa_vector(
@@ -18,9 +19,11 @@ template <typename T, int D, int V = D>
     const constant size_t& v_stride,
     const constant float& scale,
     const device bool* mask [[function_constant(has_mask)]],
-    const constant int& mask_seq_stride [[function_constant(has_mask)]],
+    const constant int& mask_kv_seq_stride [[function_constant(has_mask)]],
+    const constant int& mask_q_seq_stride [[function_constant(has_mask)]],
     const constant int& mask_head_stride [[function_constant(has_mask)]],
     uint3 tid [[threadgroup_position_in_grid]],
+    uint3 tpg [[threadgroups_per_grid]],
     uint simd_gid [[simdgroup_index_in_threadgroup]],
     uint simd_lid [[thread_index_in_simdgroup]]) {
   constexpr int BN = 32;
@@ -41,15 +44,21 @@ template <typename T, int D, int V = D>
   threadgroup U sum_exp_scores[BN];
 
   // Adjust positions
-  const int head_idx = tid.y;
+  const int head_idx = tid.x;
+  const int q_seq_idx = tid.y;
   const int kv_head_idx = head_idx / gqa_factor;
-  queries += head_idx * D + simd_lid * qk_per_thread;
+  const int o_offset = tpg.x * q_seq_idx + head_idx;
+  const int q_offset =
+      query_transposed ? o_offset : head_idx * tpg.y + q_seq_idx;
+  queries += q_offset * D + simd_lid * qk_per_thread;
   keys += kv_head_idx * k_stride + simd_gid * D + simd_lid * qk_per_thread;
   values += kv_head_idx * v_stride + simd_gid * V + simd_lid * v_per_thread;
   if (has_mask) {
-    mask += head_idx * mask_head_stride + simd_gid * mask_seq_stride;
+    mask += head_idx * mask_head_stride + simd_gid * mask_kv_seq_stride +
+        q_seq_idx * mask_q_seq_stride;
   }
-  out += head_idx * V + simd_gid * v_per_thread;
+
+  out += o_offset * V + simd_gid * v_per_thread;
 
   // Read the query and 0 the output accumulator
   for (int i = 0; i < qk_per_thread; i++) {
@@ -95,7 +104,7 @@ template <typename T, int D, int V = D>
     keys += inner_k_stride;
     values += inner_v_stride;
     if (has_mask) {
-      mask += BN * mask_seq_stride;
+      mask += BN * mask_kv_seq_stride;
     }
   }
 
@@ -142,9 +151,11 @@ template <typename T, int D, int V = D>
     const constant size_t& v_stride,
     const constant float& scale,
     const device bool* mask [[function_constant(has_mask)]],
-    const constant int& mask_seq_stride [[function_constant(has_mask)]],
+    const constant int& mask_kv_seq_stride [[function_constant(has_mask)]],
+    const constant int& mask_q_seq_stride [[function_constant(has_mask)]],
     const constant int& mask_head_stride [[function_constant(has_mask)]],
     uint3 tid [[threadgroup_position_in_grid]],
+    uint3 tpg [[threadgroups_per_grid]],
     uint simd_gid [[simdgroup_index_in_threadgroup]],
     uint simd_lid [[thread_index_in_simdgroup]]) {
   constexpr int BN = 8;
@@ -167,20 +178,26 @@ template <typename T, int D, int V = D>
 
   // Adjust positions
   const int block_idx = tid.z;
-  const int head_idx = tid.y;
+  const int head_idx = tid.x;
+  const int q_seq_idx = tid.y;
+  const int o_offset = tpg.x * q_seq_idx + head_idx;
+  const int q_offset =
+      query_transposed ? o_offset : head_idx * tpg.y + q_seq_idx;
   const int kv_head_idx = head_idx / gqa_factor;
-  queries += head_idx * D + simd_lid * qk_per_thread;
+
+  queries += q_offset * D + simd_lid * qk_per_thread;
   keys += kv_head_idx * k_stride + (block_idx * BN + simd_gid) * D +
       simd_lid * qk_per_thread;
   values += kv_head_idx * v_stride + (block_idx * BN + simd_gid) * V +
       simd_lid * v_per_thread;
-  out += head_idx * blocks * V + block_idx * V + simd_lid * v_per_thread;
+  out += o_offset * blocks * V + block_idx * V + simd_lid * v_per_thread;
   if (has_mask) {
     mask += head_idx * mask_head_stride +
-        (block_idx * BN + simd_gid) * mask_seq_stride;
+        (block_idx * BN + simd_gid) * mask_kv_seq_stride +
+        q_seq_idx * mask_q_seq_stride;
   }
-  sums += head_idx * blocks + block_idx;
-  maxs += head_idx * blocks + block_idx;
+  sums += o_offset * blocks + block_idx;
+  maxs += o_offset * blocks + block_idx;
 
   // Read the query and 0 the output accumulator
   for (int i = 0; i < qk_per_thread; i++) {
@@ -226,7 +243,7 @@ template <typename T, int D, int V = D>
     keys += blocks * inner_k_stride;
     values += blocks * inner_v_stride;
     if (has_mask) {
-      mask += BN * blocks * mask_seq_stride;
+      mask += BN * blocks * mask_kv_seq_stride;
     }
   }
 
@@ -275,6 +292,7 @@ template <typename T, int D>
     const device float* maxs [[buffer(2)]],
     device T* out [[buffer(3)]],
     uint3 tid [[threadgroup_position_in_grid]],
+    uint3 tpg [[threadgroups_per_grid]],
     uint simd_gid [[simdgroup_index_in_threadgroup]],
     uint simd_lid [[thread_index_in_simdgroup]]) {
   constexpr int BN = 32;
@@ -288,11 +306,14 @@ template <typename T, int D>
   threadgroup U outputs[BN * BD];
 
   // Adjust positions
-  const int head_idx = tid.y;
-  partials += head_idx * blocks * D + simd_gid * D + simd_lid * elem_per_thread;
-  sums += head_idx * blocks;
-  maxs += head_idx * blocks;
-  out += head_idx * D + simd_gid * elem_per_thread;
+  const int head_idx = tid.x;
+  const int q_seq_idx = tid.y;
+  const int n_heads = tpg.x;
+  const int q_offset = n_heads * q_seq_idx + head_idx;
+  partials += q_offset * blocks * D + simd_gid * D + simd_lid * elem_per_thread;
+  sums += q_offset * blocks;
+  maxs += q_offset * blocks;
+  out += q_offset * D + simd_gid * elem_per_thread;
 
   // First everybody reads the max and sum_exp
   U max_score = maxs[simd_lid];

mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.h

@@ -50,7 +50,7 @@ struct SubOp {
 struct ExpSubOp {
   template <typename T>
   METAL_FUNC static constexpr T apply(T x, T y) {
-    return fast::exp(x - y);
+    return fast::exp2(x - y);
   }
 };
 
@@ -103,17 +103,24 @@ template <
       tidl.x * BQ * params->O_strides[2]; // Seqeunce
 
   // Prepare threadgroup memory
-  constexpr short padQ = 0; // 16 / sizeof(T);
-  constexpr short padK = 0; // 16 / sizeof(T);
-  constexpr short padV = 0; // 16 / sizeof(T);
+  constexpr short padQ = 16 / sizeof(T);
+  constexpr short padK = 16 / sizeof(T);
+  constexpr short padV = 16 / sizeof(T);
 
   constexpr short LDQ_tgp = BD + padQ;
   constexpr short LDK_tgp = BK + padK;
   constexpr short LDV_tgp = BD + padV;
 
-  threadgroup T Qs[BQ * (BD + padQ)];
-  threadgroup T Ks[(BK + padK) * BD];
-  threadgroup T Vs[BK * (BD + padV)];
+  constexpr short tgp_mem_0 = (BK + padK) * (BD);
+  constexpr short tgp_mem_1 = BK * (BD + padV);
+  constexpr short tgp_mem_s = tgp_mem_0 > tgp_mem_1 ? tgp_mem_0 : tgp_mem_1;
+
+  threadgroup T Q_smem[BQ * (BD + padQ)];
+  threadgroup T KV_smem[tgp_mem_s];
+
+  threadgroup T* Qs = Q_smem;
+  threadgroup T* Ks = KV_smem;
+  threadgroup T* Vs = KV_smem;
 
   // Prepare block loaders
   using QBlockLoader = BlockLoaderT<
@@ -151,7 +158,7 @@ template <
   VBlockLoader loader_v(
       V, params->V_strides[2], Vs, simd_group_id, simd_lane_id);
 
-  TransformScale<T> ts(static_cast<T>(params->scale));
+  TransformScale<T> ts(static_cast<T>(params->scale * 1.44269504089));
 
   // Prepare MMA tiles
   constexpr short kFragSize = 8; // MMAFrag size
@@ -174,7 +181,7 @@ template <
   MMATile<AccumType, TQ, 1, MMAFrag_acc_t> Qtile;
   MMATile<AccumType, 1, TK, MMAFrag_acc_t> Ktile;
   MMATile<AccumType, TQ, TK, MMAFrag_acc_t> Stile;
-  MMATile<AccumType, TK, TD, MMAFrag_acc_t> Vtile;
+  MMATile<AccumType, 1, 1, MMAFrag_acc_t> Vtile;
   MMATile<AccumType, TQ, TD, MMAFrag_acc_t> Otile;
 
   Otile.clear();
@@ -224,11 +231,12 @@ template <
       loader_k.load_unsafe();
     }
 
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-
     // Do S = Q @ K.T
     Stile.clear();
 
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    STEEL_PRAGMA_UNROLL
     for (short dd = 0; dd < TD; dd++) {
       simdgroup_barrier(mem_flags::mem_none);
 
@@ -264,7 +272,7 @@ template <
       }
     }
 
-    simdgroup_barrier(mem_flags::mem_none);
+    threadgroup_barrier(mem_flags::mem_threadgroup);
 
     // Load V blocks
     if (!align_K && kb == (params->NK_aligned)) {
@@ -292,7 +300,7 @@ template <
     // Factor exp(rowmax(Si) - rowmax(Si-1))
     STEEL_PRAGMA_UNROLL
     for (short i = 0; i < kRowsPT; ++i) {
-      factor[i] = fast::exp(max_score[i] - new_max[i]);
+      factor[i] = fast::exp2(max_score[i] - new_max[i]);
     }
 
     // Save max for next iteration
@@ -316,12 +324,35 @@ template <
 
     // Load V into registers
     threadgroup_barrier(mem_flags::mem_threadgroup);
-    Vtile.template load<T, 1, 1, LDV_tgp, 1>(&Vs[Vs_offset]);
 
-    simdgroup_barrier(mem_flags::mem_none);
+    STEEL_PRAGMA_UNROLL
+    for (short iq = 0; iq < TQ; iq++) {
+      STEEL_PRAGMA_UNROLL
+      for (short id = 0; id < TD; id++) {
+        STEEL_PRAGMA_UNROLL
+        for (short ik = 0; ik < TK; ik++) {
+          if constexpr (BD == 128) {
+            simdgroup_barrier(mem_flags::mem_none);
+          }
 
-    // Do O = S @ V
-    tile_matmad(Otile, Stile, Vtile, Otile);
+          const short kk = ik * kFragSize;
+          const short dd = id * kFragSize;
+
+          Vtile.template load<T, 1, 1, LDV_tgp, 1>(
+              &Vs[Vs_offset + kk * LDV_tgp + dd]);
+
+          if constexpr (BD == 128) {
+            simdgroup_barrier(mem_flags::mem_none);
+          }
+
+          MMAFrag_acc_t::mma(
+              Otile.frag_at(iq, id),
+              Stile.frag_at(iq, ik),
+              Vtile.frag_at(0, 0),
+              Otile.frag_at(iq, id));
+        }
+      }
+    }
 
     // Prepare for next iteration
     loader_k.next();

mlx/backend/metal/kernels/steel/attn/mma.h

@@ -62,6 +62,12 @@ struct BaseMMAFrag<T, 8, 8> {
   typedef metal::vec<T, kElemRows> row_frag_type;
   typedef metal::vec<T, kElemCols> col_frag_type;
 
+  template <typename U>
+  using dtype_mat_t = typename metal::simdgroup_matrix<U, kFragRows, kFragCols>;
+
+  template <typename U>
+  using dtype_frag_t = typename metal::vec<U, kElemsPerFrag>;
+
   METAL_FUNC static constexpr short2 get_coord(ushort simd_lane_id
                                                [[thread_index_in_simdgroup]]) {
     const short qid = simd_lane_id / 4;
@@ -158,30 +164,32 @@ struct BaseMMAFrag<T, 8, 8> {
     }
   }
 
+  template <typename Atype, typename Btype, typename Ctype>
   METAL_FUNC static constexpr void mma(
       thread frag_type& D,
-      thread frag_type& A,
-      thread frag_type& B,
-      thread frag_type& C) {
+      thread dtype_frag_t<Atype>& A,
+      thread dtype_frag_t<Btype>& B,
+      thread dtype_frag_t<Ctype>& C) {
     mat_type D_mat;
-    mat_type A_mat;
-    mat_type B_mat;
-    mat_type C_mat;
+    dtype_mat_t<Atype> A_mat;
+    dtype_mat_t<Btype> B_mat;
+    dtype_mat_t<Ctype> C_mat;
 
-    reinterpret_cast<thread frag_type&>(A_mat.thread_elements()) = A;
-    reinterpret_cast<thread frag_type&>(B_mat.thread_elements()) = B;
-    reinterpret_cast<thread frag_type&>(C_mat.thread_elements()) = C;
+    reinterpret_cast<thread dtype_frag_t<Atype>&>(A_mat.thread_elements()) = A;
+    reinterpret_cast<thread dtype_frag_t<Btype>&>(B_mat.thread_elements()) = B;
+    reinterpret_cast<thread dtype_frag_t<Ctype>&>(C_mat.thread_elements()) = C;
 
     mma(D_mat, A_mat, B_mat, C_mat);
 
     D = reinterpret_cast<thread frag_type&>(D_mat.thread_elements());
   }
 
+  template <typename Atype, typename Btype, typename Ctype>
   METAL_FUNC static constexpr void mma(
       thread mat_type& D,
-      thread mat_type& A,
-      thread mat_type& B,
-      thread mat_type& C) {
+      thread dtype_mat_t<Atype>& A,
+      thread dtype_mat_t<Btype>& B,
+      thread dtype_mat_t<Ctype>& C) {
     simdgroup_multiply_accumulate(D, A, B, C);
   }
 
@@ -242,7 +250,7 @@ struct MMATile {
   typedef typename MMAFrag_t::mat_type mat_type;
   typedef typename MMAFrag_t::frag_type frag_type;
 
-  frag_type val_frags[kNumFrags] = {frag_type(0)};
+  frag_type val_frags[kNumFrags]; // = {frag_type(0)};
 
   METAL_FUNC MMATile() thread {}
 
@@ -409,24 +417,37 @@ struct MMATile {
   }
 };
 
-template <typename T, typename U, int M, int N, int K>
+template <
+    typename Dtype,
+    typename Atype,
+    typename Btype,
+    typename Ctype,
+    int M,
+    int N,
+    int K,
+    class MMAFragD,
+    class MMAFragA,
+    class MMAFragB,
+    class MMAFragC>
 METAL_FUNC void tile_matmad(
-    thread MMATile<T, M, N>& D,
-    thread MMATile<U, M, K>& A,
-    thread MMATile<U, K, N>& B,
-    thread MMATile<T, M, N>& C) {
+    thread MMATile<Dtype, M, N, MMAFragD>& D,
+    thread MMATile<Atype, M, K, MMAFragA>& A,
+    thread MMATile<Btype, K, N, MMAFragB>& B,
+    thread MMATile<Ctype, M, N, MMAFragC>& C) {
   STEEL_PRAGMA_UNROLL
-  for (short k = 0; k < K; ++k) {
+  for (short m = 0; m < M; ++m) {
     STEEL_PRAGMA_UNROLL
-    for (short m = 0; m < M; ++m) {
+    for (short n = 0; n < N; ++n) {
+      short m_serp = m; //(n % 2) ? (M - 1 - m) : m;
+      short n_serp = (m % 2) ? (N - 1 - n) : n;
+
       STEEL_PRAGMA_UNROLL
-      for (short n = 0; n < N; ++n) {
-        short n_serp = (m % 2) ? (N - 1 - n) : n;
-        MMATile<T, M, N>::MMAFrag_t::mma(
-            D.frag_at(m, n_serp),
-            A.frag_at(m, k),
+      for (short k = 0; k < K; ++k) {
+        MMAFragD::mma(
+            D.frag_at(m_serp, n_serp),
+            A.frag_at(m_serp, k),
             B.frag_at(k, n_serp),
-            C.frag_at(m, n_serp));
+            C.frag_at(m_serp, n_serp));
       }
     }
   }

mlx/backend/metal/metal.h

@@ -82,7 +82,7 @@ void start_capture(std::string path = "");
 void stop_capture();
 
 /** Get information about the GPU and system settings. */
-std::unordered_map<std::string, std::variant<std::string, size_t>>
+const std::unordered_map<std::string, std::variant<std::string, size_t>>&
 device_info();
 
 } // namespace mlx::core::metal

mlx/backend/metal/normalization.cpp

@@ -77,7 +77,7 @@ void RMSNorm::eval_gpu(
       group_dims = MTL::Size(threadgroup_size, 1, 1);
     }
 
-    uint32_t w_stride = w.strides()[0];
+    uint32_t w_stride = (w.ndim() == 1) ? w.strides()[0] : 0;
     compute_encoder.set_compute_pipeline_state(kernel);
     compute_encoder.set_input_array(
         x.data_shared_ptr() == nullptr ? out : x, 0);
@@ -101,20 +101,16 @@ void RMSNormVJP::eval_gpu(
   // Ensure row contiguity. We could relax this step by checking that the array
   // is contiguous (no broadcasts or holes) and that the input strides are the
   // same as the cotangent strides but for now this is simpler.
-  std::vector<array> copies;
-  auto check_input = [&copies, &s](const array& x) -> const array& {
+  auto check_input = [&d, &s](const array& x) -> array {
     if (x.flags().row_contiguous) {
       return x;
     }
-    // Make sure we 'll only ever allocate once. The point of that goes beyond
-    // the minor optimization. We need to ensure that there will be no
-    // reallocation such that the references won't change when we
-    // push_back(...). So tl;dr 3 possible copies x, g and gw_temp.
-    copies.reserve(3);
 
-    copies.push_back(array(x.shape(), x.dtype(), nullptr, {}));
-    copy_gpu(x, copies.back(), CopyType::General, s);
-    return copies.back();
+    array x_copy(x.shape(), x.dtype(), nullptr, {});
+    copy_gpu(x, x_copy, CopyType::General, s);
+    d.add_temporary(x_copy, s.index);
+
+    return x_copy;
   };
   const array& x = check_input(inputs[0]);
   const array& w = inputs[1];
@@ -122,6 +118,9 @@ void RMSNormVJP::eval_gpu(
   array& gx = outputs[0];
   array& gw = outputs[1];
 
+  // Check whether we had a weight
+  bool has_w = w.ndim() != 0;
+
   // Allocate space for the outputs
   bool x_in_gx = false;
   bool g_in_gx = false;
@@ -140,15 +139,18 @@ void RMSNormVJP::eval_gpu(
 
   // Allocate the gradient accumulator gw and a temporary to store the
   // gradients before they are accumulated.
-  array gw_temp({n_rows, x.shape().back()}, gw.dtype(), nullptr, {});
+  array gw_temp =
+      (has_w) ? array({n_rows, x.shape().back()}, gw.dtype(), nullptr, {}) : w;
   bool g_in_gw = false;
-  if (!g_in_gx && g.is_donatable()) {
-    gw_temp.move_shared_buffer(g);
-    g_in_gw = true;
-  } else {
-    gw_temp.set_data(allocator::malloc_or_wait(gw_temp.nbytes()));
+  if (has_w) {
+    if (!g_in_gx && g.is_donatable()) {
+      gw_temp.move_shared_buffer(g);
+      g_in_gw = true;
+    } else {
+      gw_temp.set_data(allocator::malloc_or_wait(gw_temp.nbytes()));
+      d.add_temporary(gw_temp, s.index);
+    }
   }
-  copies.push_back(gw_temp);
   gw.set_data(allocator::malloc_or_wait(gw.nbytes()));
 
   const int simd_size = 32;
@@ -159,9 +161,15 @@ void RMSNormVJP::eval_gpu(
     op_name += "_looped";
   }
   op_name += type_to_name(gx);
+
+  std::string hash_name = op_name + ((has_w) ? "_w" : "_now");
+  metal::MTLFCList func_consts = {
+      {&has_w, MTL::DataType::DataTypeBool, 20},
+  };
+
   auto& compute_encoder = d.get_command_encoder(s.index);
   {
-    auto kernel = d.get_kernel(op_name);
+    auto kernel = d.get_kernel(op_name, "mlx", hash_name, func_consts);
 
     MTL::Size grid_dims, group_dims;
     if (axis_size <= looped_limit) {
@@ -179,7 +187,7 @@ void RMSNormVJP::eval_gpu(
       group_dims = MTL::Size(threadgroup_size, 1, 1);
     }
 
-    uint32_t w_stride = w.strides()[0];
+    uint32_t w_stride = (w.ndim() == 1) ? w.strides()[0] : 0;
     compute_encoder.set_compute_pipeline_state(kernel);
     compute_encoder.set_input_array(x_in_gx ? gx : x, 0);
     compute_encoder.set_input_array(w, 1);
@@ -192,12 +200,12 @@ void RMSNormVJP::eval_gpu(
     compute_encoder.dispatch_threads(grid_dims, group_dims);
   }
 
-  ReductionPlan plan(
-      ReductionOpType::ContiguousStridedReduce, {n_rows}, {axis_size});
-  strided_reduce_general_dispatch(
-      gw_temp, gw, "sum", plan, {0}, compute_encoder, d, s);
-
-  d.add_temporaries(std::move(copies), s.index);
+  if (has_w) {
+    ReductionPlan plan(
+        ReductionOpType::ContiguousStridedReduce, {n_rows}, {axis_size});
+    strided_reduce_general_dispatch(
+        gw_temp, gw, "sum", plan, {0}, compute_encoder, d, s);
+  }
 }
 
 void LayerNorm::eval_gpu(
@@ -295,20 +303,16 @@ void LayerNormVJP::eval_gpu(
   // Ensure row contiguity. We could relax this step by checking that the array
   // is contiguous (no broadcasts or holes) and that the input strides are the
   // same as the cotangent strides but for now this is simpler.
-  std::vector<array> copies;
-  auto check_input = [&copies, &s](const array& x) -> const array& {
+  auto check_input = [&d, &s](const array& x) -> array {
     if (x.flags().row_contiguous) {
       return x;
     }
-    // Make sure we 'll only ever allocate once. The point of that goes beyond
-    // the minor optimization. We need to ensure that there will be no
-    // reallocation such that the references won't change when we
-    // push_back(...). So tl;dr 3 possible copies x, g and gw_temp.
-    copies.reserve(3);
 
-    copies.push_back(array(x.shape(), x.dtype(), nullptr, {}));
-    copy_gpu(x, copies.back(), CopyType::General, s);
-    return copies.back();
+    array x_copy(x.shape(), x.dtype(), nullptr, {});
+    copy_gpu(x, x_copy, CopyType::General, s);
+    d.add_temporary(x_copy, s.index);
+
+    return x_copy;
   };
   const array& x = check_input(inputs[0]);
   const array& w = inputs[1];
@@ -318,6 +322,9 @@ void LayerNormVJP::eval_gpu(
   array& gw = outputs[1];
   array& gb = outputs[2];
 
+  // Check whether we had a weight
+  bool has_w = w.ndim() != 0;
+
   // Allocate space for the outputs
   bool x_in_gx = false;
   bool g_in_gx = false;
@@ -336,15 +343,18 @@ void LayerNormVJP::eval_gpu(
 
   // Allocate a temporary to store the gradients for w and allocate the output
   // gradient accumulators.
-  array gw_temp({n_rows, x.shape().back()}, gw.dtype(), nullptr, {});
+  array gw_temp =
+      (has_w) ? array({n_rows, x.shape().back()}, gw.dtype(), nullptr, {}) : w;
   bool g_in_gw = false;
-  if (!g_in_gx && g.is_donatable()) {
-    gw_temp.move_shared_buffer(g);
-    g_in_gw = true;
-  } else {
-    gw_temp.set_data(allocator::malloc_or_wait(gw_temp.nbytes()));
+  if (has_w) {
+    if (!g_in_gx && g.is_donatable()) {
+      gw_temp.move_shared_buffer(g);
+      g_in_gw = true;
+    } else {
+      gw_temp.set_data(allocator::malloc_or_wait(gw_temp.nbytes()));
+    }
+    d.add_temporary(gw_temp, s.index);
   }
-  copies.push_back(gw_temp);
   gw.set_data(allocator::malloc_or_wait(gw.nbytes()));
   gb.set_data(allocator::malloc_or_wait(gb.nbytes()));
 
@@ -372,8 +382,14 @@ void LayerNormVJP::eval_gpu(
     op_name += "_looped";
   }
   op_name += type_to_name(gx);
+
+  std::string hash_name = op_name + ((has_w) ? "_w" : "_now");
+  metal::MTLFCList func_consts = {
+      {&has_w, MTL::DataType::DataTypeBool, 20},
+  };
+
   {
-    auto kernel = d.get_kernel(op_name);
+    auto kernel = d.get_kernel(op_name, "mlx", hash_name, func_consts);
 
     MTL::Size grid_dims, group_dims;
     if (axis_size <= looped_limit) {
@@ -404,14 +420,12 @@ void LayerNormVJP::eval_gpu(
     compute_encoder.dispatch_threads(grid_dims, group_dims);
   }
 
-  if (gw.ndim() == 1 && gw.size() == axis_size) {
+  if (has_w) {
     ReductionPlan plan(
         ReductionOpType::ContiguousStridedReduce, {n_rows}, {axis_size});
     strided_reduce_general_dispatch(
         gw_temp, gw, "sum", plan, {0}, compute_encoder, d, s);
   }
-
-  d.add_temporaries(std::move(copies), s.index);
 }
 
 } // namespace mlx::core::fast

mlx/backend/metal/rope.cpp

@@ -25,6 +25,10 @@ void RoPE::eval_gpu(
   size_t out_strides[3];
   bool donated = false;
   int ndim = in.ndim();
+  int dispatch_ndim = in.ndim();
+  while (in.shape(-dispatch_ndim) == 1 && dispatch_ndim > 3) {
+    dispatch_ndim--;
+  }
   size_t mat_size = in.shape(-2) * in.shape(-1);
   if (dims_ < in.shape(-1)) {
     donated = true;
@@ -44,12 +48,12 @@ void RoPE::eval_gpu(
     strides[0] = mat_size;
     strides[1] = in.strides()[ndim - 2];
     strides[2] = in.strides()[ndim - 1];
-  } else if (ndim == 3) {
+  } else if (dispatch_ndim == 3) {
     // Handle non-contiguous 3D inputs
     out.set_data(allocator::malloc_or_wait(out.nbytes()));
-    strides[0] = in.strides()[0];
-    strides[1] = in.strides()[1];
-    strides[2] = in.strides()[2];
+    strides[0] = in.strides()[ndim - 3];
+    strides[1] = in.strides()[ndim - 2];
+    strides[2] = in.strides()[ndim - 1];
   } else {
     // Copy non-contiguous > 3D inputs into the output and treat
     // input as donated

mlx/backend/metal/scaled_dot_product_attention.cpp

@@ -134,14 +134,17 @@ void sdpa_vector(
   size_t k_stride = k.strides()[1];
   size_t v_stride = v.strides()[1];
   MTL::Size group_dims(1024, 1, 1);
-  MTL::Size grid_dims(1, B, 1);
+  MTL::Size grid_dims(B, q.shape(2), 1);
 
   bool has_mask = mask.has_value();
+  bool query_transposed = !q.flags().row_contiguous;
   metal::MTLFCList func_consts = {
       {&has_mask, MTL::DataType::DataTypeBool, 20},
+      {&query_transposed, MTL::DataType::DataTypeBool, 21},
   };
   std::string hash_name = kname;
   hash_name += has_mask ? "_mask" : "_nomask";
+  hash_name += query_transposed ? "_qt" : "_qnt";
 
   // Get the kernel
   auto& compute_encoder = d.get_command_encoder(s.index);
@@ -161,10 +164,14 @@ void sdpa_vector(
   if (has_mask) {
     auto& m = *mask;
     compute_encoder.set_input_array(m, 9);
-    int32_t seq_stride = m.ndim() >= 1 ? m.strides().back() : 0;
-    int32_t head_stride = m.ndim() >= 3 ? *(m.strides().end() - 3) : 0;
-    compute_encoder.set_bytes(seq_stride, 10);
-    compute_encoder.set_bytes(head_stride, 11);
+    auto nd = m.ndim();
+    int32_t kv_seq_stride =
+        nd >= 1 && m.shape(-1) > 1 ? m.strides()[nd - 1] : 0;
+    int32_t q_seq_stride = nd >= 2 && m.shape(-2) > 1 ? m.strides()[nd - 2] : 0;
+    int32_t head_stride = nd >= 3 && m.shape(-3) > 1 ? m.strides()[nd - 3] : 0;
+    compute_encoder.set_bytes(kv_seq_stride, 10);
+    compute_encoder.set_bytes(q_seq_stride, 11);
+    compute_encoder.set_bytes(head_stride, 12);
   }
 
   // Launch
@@ -198,7 +205,7 @@ void sdpa_vector_2pass(
   auto k_stride = k.strides()[1];
   auto v_stride = v.strides()[1];
   MTL::Size group_dims(8 * 32, 1, 1);
-  MTL::Size grid_dims(1, B, blocks);
+  MTL::Size grid_dims(B, q.shape(2), blocks);
 
   // Allocate the intermediates
   Shape intermediate_shape;
@@ -219,11 +226,14 @@ void sdpa_vector_2pass(
   d.add_temporary(maxs, s.index);
 
   bool has_mask = mask.has_value();
+  bool query_transposed = !q.flags().row_contiguous;
   metal::MTLFCList func_consts = {
       {&has_mask, MTL::DataType::DataTypeBool, 20},
+      {&query_transposed, MTL::DataType::DataTypeBool, 21},
   };
   std::string hash_name = kname;
   hash_name += has_mask ? "_mask" : "_nomask";
+  hash_name += query_transposed ? "_qt" : "_qnt";
 
   // Get the kernel
   auto& compute_encoder = d.get_command_encoder(s.index);
@@ -246,10 +256,14 @@ void sdpa_vector_2pass(
   if (has_mask) {
     auto& m = *mask;
     compute_encoder.set_input_array(m, 11);
-    int32_t seq_stride = m.ndim() >= 1 ? m.strides().back() : 0;
-    int32_t head_stride = m.ndim() >= 3 ? *(m.strides().end() - 3) : 0;
-    compute_encoder.set_bytes(seq_stride, 12);
-    compute_encoder.set_bytes(head_stride, 13);
+    auto nd = m.ndim();
+    int32_t kv_seq_stride =
+        nd >= 1 && m.shape(-1) > 1 ? m.strides()[nd - 1] : 0;
+    int32_t q_seq_stride = nd >= 2 && m.shape(-2) > 1 ? m.strides()[nd - 2] : 0;
+    int32_t head_stride = nd >= 3 && m.shape(-3) > 1 ? m.strides()[nd - 3] : 0;
+    compute_encoder.set_bytes(kv_seq_stride, 12);
+    compute_encoder.set_bytes(q_seq_stride, 13);
+    compute_encoder.set_bytes(head_stride, 14);
   }
 
   // Launch
@@ -274,7 +288,7 @@ void sdpa_vector_2pass(
 
   // Launch
   group_dims = MTL::Size(1024, 1, 1);
-  grid_dims = MTL::Size(1, B, 1);
+  grid_dims = MTL::Size(B, q.shape(2), 1);
   compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
 }
 
@@ -301,16 +315,23 @@ void ScaledDotProductAttention::eval_gpu(
     if (!predicate(arr)) {
       array arr_copy(arr.shape(), arr.dtype(), nullptr, {});
       copy_gpu(arr, arr_copy, CopyType::General, s);
-      copies.push_back(arr_copy);
+      copies.push_back(std::move(arr_copy));
       return copies.back();
     } else {
       return arr;
     }
   };
 
-  // Checks if arr is fully row contiguous
-  auto is_contiguous = [](const array& arr) {
-    return arr.flags().row_contiguous;
+  // Checks if arr is row contiguous or the sequence and head dimension are
+  // transposed
+  auto is_contiguous_or_head_seq_transposed = [](const array& arr) {
+    if (arr.flags().row_contiguous) {
+      return true;
+    }
+    auto& strides = arr.strides();
+    auto& shape = arr.shape();
+    return (strides[3] == 1) && (strides[2] == shape[3] * shape[1]) &&
+        (strides[1] == shape[3]) && (strides[0] == strides[2] * shape[2]);
   };
 
   // Returns true if the array is row contiguous except the sequence length
@@ -328,18 +349,30 @@ void ScaledDotProductAttention::eval_gpu(
   };
 
   // We are in vector mode ie single query
-  if (q_pre.shape(2) == 1) {
-    const auto& q = copy_unless(is_contiguous, q_pre);
-    // 1, heads, seq_len, head_dim
-    // mask [1, query_heads, 1, seq_len]
+  if (q_pre.shape(2) <= 8) {
+    const auto& q = copy_unless(is_contiguous_or_head_seq_transposed, q_pre);
     const auto& k = copy_unless(is_contiguous_except_seq_len, k_pre);
     const auto& v = copy_unless(is_contiguous_except_seq_len, v_pre);
 
     // Donate the query if possible
-    if (q.is_donatable() && q.size() == o.size()) {
+    if (q.is_donatable() && (q.shape(2) == 1 || !q.flags().row_contiguous) &&
+        q.size() == o.size()) {
       o.move_shared_buffer(q);
     } else {
-      o.set_data(allocator::malloc_or_wait(o.nbytes()));
+      if (o.shape(2) == 1) {
+        o.set_data(allocator::malloc_or_wait(o.nbytes()));
+      } else {
+        auto strides = o.strides();
+        strides[2] = o.shape(1) * o.shape(3);
+        strides[1] = o.shape(3);
+        auto flags = q.flags();
+        flags.row_contiguous = q.shape(1) == 1;
+        o.set_data(
+            allocator::malloc_or_wait(o.nbytes()),
+            o.size(),
+            std::move(strides),
+            flags);
+      }
     }
 
     auto mask =

mlx/backend/metal/scan.cpp

@@ -17,10 +17,10 @@ void Scan::eval_gpu(const std::vector<array>& inputs, array& out) {
   auto& s = stream();
   auto& d = metal::device(s.device);
 
-  std::vector<array> copies;
+  bool donate = inputs[0].is_donatable();
   auto in = inputs[0];
   if (in.flags().contiguous && in.strides()[axis_] != 0) {
-    if (in.is_donatable() && in.itemsize() == out.itemsize()) {
+    if (donate && in.itemsize() == out.itemsize()) {
       out.move_shared_buffer(in);
     } else {
       out.set_data(
@@ -32,8 +32,7 @@ void Scan::eval_gpu(const std::vector<array>& inputs, array& out) {
   } else {
     array arr_copy(in.shape(), in.dtype(), nullptr, {});
     copy_gpu(in, arr_copy, CopyType::General, s);
-    copies.push_back(arr_copy);
-    in = arr_copy;
+    in = std::move(arr_copy);
     out.move_shared_buffer(in);
   }
 
@@ -127,8 +126,6 @@ void Scan::eval_gpu(const std::vector<array>& inputs, array& out) {
     MTL::Size group_dims(thread_group_size, 1, 1);
     compute_encoder.dispatch_threads(grid_dims, group_dims);
   }
-
-  d.add_temporaries(std::move(copies), s.index);
 }
 
 } // namespace mlx::core

mlx/backend/no_metal/metal.cpp

@@ -54,7 +54,7 @@ void start_capture(std::string) {}
 void stop_capture() {}
 void clear_cache() {}
 
-std::unordered_map<std::string, std::variant<std::string, size_t>>
+const std::unordered_map<std::string, std::variant<std::string, size_t>>&
 device_info() {
   throw std::runtime_error(
       "[metal::device_info] Cannot get device info without metal backend");

mlx/compile.cpp

@@ -15,6 +15,7 @@
 namespace mlx::core {
 
 constexpr int max_compile_depth = 11;
+constexpr int max_compile_arrays = 24;
 
 bool is_unary(const Primitive& p) {
   return (
@@ -570,6 +571,7 @@ void compile_fuse(
 
     std::function<void(const array&, int, const Stream&, const Shape&)> recurse;
     std::unordered_set<uintptr_t> cache;
+    std::unordered_set<uintptr_t> input_set;
     recurse = [&](const array& a,
                   int depth,
                   const Stream& s,
@@ -587,6 +589,8 @@ void compile_fuse(
       if (depth >= max_compile_depth || !a.has_primitive() ||
           a.primitive().stream() != s || !is_fusable(a.primitive()) ||
           (output_map.find(a.id()) != output_map.end() && a.shape() != shape)) {
+        // Possible input
+        input_set.insert(a.id());
         return;
       }
 
@@ -607,9 +611,20 @@ void compile_fuse(
       // Arrays with a mix of parents outside the compilable section
       // are not fusable
       if (!all_parents_in) {
+        // Possible input
+        input_set.insert(a.id());
         return;
       }
 
+      if (output_map.find(a.id()) != output_map.end()) {
+        input_set.insert(a.id());
+      } else {
+        // Not an input anymore since fusing it
+        input_set.erase(a.id());
+      }
+      if (input_set.size() >= max_compile_arrays) {
+        return;
+      }
       cache.insert({a.id()});
 
       for (auto& in : a.inputs()) {
@@ -630,7 +645,7 @@ void compile_fuse(
 
     // Recurse a second time to build the tape in the right
     // order and collect the inputs
-    std::unordered_set<uintptr_t> input_set;
+    input_set.clear();
     std::vector<array> inputs;
     std::vector<array> fused_tape;
     std::unordered_set<uintptr_t> tape_set;

mlx/distributed/ring/ring.cpp

@@ -1,15 +1,20 @@
 // Copyright © 2024 Apple Inc.
 
 #include <arpa/inet.h>
+#include <fcntl.h>
 #include <netdb.h>
+#include <netinet/tcp.h>
 #include <sys/socket.h>
 #include <unistd.h>
 
 #include <chrono>
 #include <fstream>
+#include <future>
 #include <iostream>
+#include <list>
 #include <sstream>
 #include <thread>
+#include <unordered_map>
 
 #include <json.hpp>
 
@@ -18,6 +23,10 @@
 #include "mlx/distributed/distributed_impl.h"
 #include "mlx/threadpool.h"
 
+#ifndef SOL_TCP
+#define SOL_TCP IPPROTO_TCP
+#endif
+
 #define SWITCH_TYPE(x, ...)  \
   switch ((x).dtype()) {     \
     case bool_: {            \
@@ -80,53 +89,17 @@
 
 namespace mlx::core::distributed::ring {
 
-constexpr const size_t PACKET_SIZE = 262144;
+constexpr const size_t ALL_SUM_SIZE = 8 * 1024 * 1024;
+constexpr const size_t ALL_SUM_BUFFERS = 2;
 constexpr const int CONN_ATTEMPTS = 5;
 constexpr const int CONN_WAIT = 1000;
 
 using GroupImpl = mlx::core::distributed::detail::GroupImpl;
 using json = nlohmann::json;
+using namespace std::chrono_literals;
 
 namespace {
 
-class Barrier {
- public:
-  explicit Barrier(int n_threads)
-      : n_threads_(n_threads), count_(0), flag_(false) {}
-
-  void arrive_and_wait() {
-    std::unique_lock<std::mutex> lock(mtx_);
-
-    // Keep the flag that marks the current use of the barrier. The next use is
-    // going to have this flag flipped.
-    bool initial_flag = flag_;
-
-    // Increment the count
-    count_++;
-
-    // We are the last thread to arrive so reset the count, change the flag and
-    // notify everybody.
-    if (count_ == n_threads_) {
-      count_ = 0;
-      flag_ = !flag_;
-      cv_.notify_all();
-    }
-
-    // Wait for the rest to arrive
-    else {
-      cv_.wait(lock, [this, initial_flag]() { return initial_flag != flag_; });
-    }
-  }
-
- private:
-  std::mutex mtx_;
-  std::condition_variable cv_;
-  int n_threads_;
-
-  int count_;
-  bool flag_; // we need this for sequential use of the barrier
-};
-
 template <typename T>
 void log(std::ostream& os, T first) {
   os << first << std::endl;
@@ -151,6 +124,177 @@ decltype(T() * U()) ceildiv(T a, U b) {
   return (a + b - 1) / b;
 }
 
+class SocketThread {
+ public:
+  SocketThread(int fd) : fd_(fd), stop_(false) {
+    worker_ = std::thread(&SocketThread::worker, this);
+    int flags = fcntl(fd, F_GETFL, 0);
+    fcntl(fd, F_SETFL, flags | O_NONBLOCK);
+  }
+  ~SocketThread() {
+    stop_ = true;
+    condition_.notify_all();
+    worker_.join();
+    int flags = fcntl(fd_, F_GETFL, 0);
+    fcntl(fd_, F_SETFL, flags & ~O_NONBLOCK);
+  }
+
+  template <typename T>
+  std::future<void> send(T* buffer, size_t size) {
+    return send_impl(reinterpret_cast<char*>(buffer), size * sizeof(T));
+  }
+
+  template <typename T>
+  std::future<void> recv(T* buffer, size_t size) {
+    return recv_impl(reinterpret_cast<char*>(buffer), size * sizeof(T));
+  }
+
+ private:
+  struct SocketTask {
+    SocketTask(void* b, size_t s, std::promise<void>&& p)
+        : buffer(b), size(s), promise(std::move(p)) {}
+    SocketTask(SocketTask&& t)
+        : buffer(t.buffer), size(t.size), promise(std::move(t.promise)) {}
+    void* buffer;
+    size_t size;
+    std::promise<void> promise;
+  };
+
+  std::future<void> send_impl(char* buffer, size_t size) {
+    std::promise<void> send_completed_promise;
+    auto send_completed_future = send_completed_promise.get_future();
+    if (size == 0) {
+      send_completed_promise.set_value();
+      return send_completed_future;
+    }
+
+    {
+      std::unique_lock lock(queue_mutex_);
+      sends_.emplace_back(
+          SocketTask(buffer, size, std::move(send_completed_promise)));
+    }
+    condition_.notify_one();
+    return send_completed_future;
+  }
+
+  std::future<void> recv_impl(char* buffer, size_t size) {
+    std::promise<void> recv_completed_promise;
+    auto recv_completed_future = recv_completed_promise.get_future();
+    if (size == 0) {
+      recv_completed_promise.set_value();
+      return recv_completed_future;
+    }
+
+    {
+      std::unique_lock lock(queue_mutex_);
+      recvs_.emplace_back(
+          SocketTask(buffer, size, std::move(recv_completed_promise)));
+    }
+    condition_.notify_one();
+    return recv_completed_future;
+  }
+
+  bool have_tasks() {
+    return !(sends_.empty() && recvs_.empty());
+  }
+
+  void worker() {
+    int error_count = 0;
+    bool delete_recv = false;
+    bool delete_send = false;
+    while (true) {
+      {
+        std::unique_lock lock(queue_mutex_);
+
+        if (delete_recv) {
+          recvs_.front().promise.set_value();
+          recvs_.pop_front();
+          delete_recv = false;
+        }
+        if (delete_send) {
+          sends_.front().promise.set_value();
+          sends_.pop_front();
+          delete_send = false;
+        }
+
+        if (stop_) {
+          return;
+        }
+
+        if (!have_tasks()) {
+          condition_.wait(lock, [this] { return stop_ || have_tasks(); });
+          if (stop_) {
+            return;
+          }
+        }
+      }
+
+      if (!recvs_.empty()) {
+        auto& task = recvs_.front();
+        ssize_t r = ::recv(fd_, task.buffer, task.size, 0);
+        if (r > 0) {
+          task.buffer = static_cast<char*>(task.buffer) + r;
+          task.size -= r;
+          delete_recv = task.size == 0;
+          error_count = 0;
+        } else if (errno != EAGAIN) {
+          error_count++;
+          log_info(
+              true, "Receiving from socket", fd_, "failed with errno", errno);
+        }
+      }
+      if (!sends_.empty()) {
+        auto& task = sends_.front();
+        ssize_t r = ::send(fd_, task.buffer, task.size, 0);
+        if (r > 0) {
+          task.buffer = static_cast<char*>(task.buffer) + r;
+          task.size -= r;
+          delete_send = task.size == 0;
+          error_count = 0;
+        } else if (errno != EAGAIN) {
+          error_count++;
+          log_info(true, "Sending to socket", fd_, "failed with errno", errno);
+        }
+      }
+
+      if (error_count >= 10) {
+        log_info(true, "Too many send/recv errors. Aborting...");
+        return;
+      }
+    }
+  }
+
+  int fd_;
+  bool stop_;
+  std::thread worker_;
+  std::mutex queue_mutex_;
+  std::condition_variable condition_;
+  std::list<SocketTask> sends_;
+  std::list<SocketTask> recvs_;
+};
+
+class CommunicationThreads {
+ public:
+  void add(const std::vector<int>& sockets) {
+    for (int sock : sockets) {
+      threads_.emplace(sock, sock);
+    }
+  }
+
+  template <typename T>
+  std::future<void> send(int socket, T* buffer, size_t size) {
+    return threads_.at(socket).send<T>(buffer, size);
+  }
+
+  template <typename T>
+  std::future<void> recv(int socket, T* buffer, size_t size) {
+    return threads_.at(socket).recv<T>(buffer, size);
+  }
+
+ private:
+  std::unordered_map<int, SocketThread> threads_;
+};
+
 struct address_t {
   sockaddr_storage addr;
   socklen_t len;
@@ -378,140 +522,6 @@ void sum_inplace(const T* input, T* output, size_t N) {
   }
 }
 
-template <typename T>
-void _send(int sock, T* data, size_t start, size_t stop) {
-  if (stop <= start) {
-    return;
-  }
-  data += start;
-  size_t len = (stop - start) * sizeof(T);
-  const char* buffer = (const char*)data;
-  while (len > 0) {
-    ssize_t r = send(sock, buffer, len, 0);
-    if (r <= 0) {
-      std::ostringstream msg;
-      msg << "Send of " << len << " bytes failed (errno: " << errno << ")";
-      throw std::runtime_error(msg.str());
-    }
-    buffer += r;
-    len -= r;
-  }
-}
-
-template <typename T>
-void _recv(int sock, T* data, size_t start, size_t stop) {
-  if (stop <= start) {
-    return;
-  }
-  data += start;
-  size_t len = (stop - start) * sizeof(T);
-  char* buffer = (char*)data;
-  while (len > 0) {
-    ssize_t r = recv(sock, buffer, len, 0);
-    if (r <= 0) {
-      std::ostringstream msg;
-      msg << "Recv of " << len << " bytes failed (errno: " << errno << ")";
-      throw std::runtime_error(msg.str());
-    }
-    buffer += r;
-    len -= r;
-  }
-}
-
-template <typename T>
-void _recv_sum(int sock, T* data, size_t start, size_t stop) {
-  if (stop <= start) {
-    return;
-  }
-  data += start;
-  char buffer[PACKET_SIZE];
-  size_t len = (stop - start) * sizeof(T);
-  while (len > 0) {
-    ssize_t r = 0;
-    do {
-      ssize_t partial_r =
-          recv(sock, buffer + r, std::min(len, PACKET_SIZE) - r, 0);
-      if (partial_r <= 0) {
-        std::ostringstream msg;
-        msg << "Recv of " << len << " bytes failed (errno: " << errno << ")";
-        throw std::runtime_error(msg.str());
-      }
-      r += partial_r;
-    } while (r % sizeof(T));
-    sum_inplace((const T*)buffer, data, r / sizeof(T));
-    data += r / sizeof(T);
-    len -= r;
-  }
-}
-
-template <typename T>
-void ring_send(
-    Barrier& barrier,
-    int socket,
-    int rank,
-    int size,
-    T* data,
-    size_t data_size,
-    int direction = -1) {
-  // We split the data into `size_` segments of size `segment_size`
-  size_t segment_size = ceildiv(data_size, size);
-
-  // Initial segment
-  int segment = rank;
-
-  // 1st send
-  for (int i = 0; i < size - 1; i++) {
-    size_t start = segment * segment_size;
-    size_t stop = std::min((segment + 1) * segment_size, data_size);
-    _send<T>(socket, data, start, stop);
-    barrier.arrive_and_wait();
-    segment = (segment + size + direction) % size;
-  }
-
-  // 2nd send
-  for (int i = 0; i < size - 1; i++) {
-    size_t start = segment * segment_size;
-    size_t stop = std::min((segment + 1) * segment_size, data_size);
-    _send<T>(socket, data, start, stop);
-    barrier.arrive_and_wait();
-    segment = (segment + size + direction) % size;
-  }
-}
-
-template <typename T>
-void ring_recv_sum(
-    Barrier& barrier,
-    int socket,
-    int rank,
-    int size,
-    T* data,
-    size_t data_size,
-    int direction = -1) {
-  // We split the data into `size_` segments of size `segment_size`
-  size_t segment_size = ceildiv(data_size, size);
-
-  // Initial segment
-  int segment = (rank + size + direction) % size;
-
-  // Recv sum
-  for (int i = 0; i < size - 1; i++) {
-    size_t start = segment * segment_size;
-    size_t stop = std::min((segment + 1) * segment_size, data_size);
-    _recv_sum<T>(socket, data, start, stop);
-    barrier.arrive_and_wait();
-    segment = (segment + size + direction) % size;
-  }
-
-  // Recv
-  for (int i = 0; i < size - 1; i++) {
-    size_t start = segment * segment_size;
-    size_t stop = std::min((segment + 1) * segment_size, data_size);
-    _recv<T>(socket, data, start, stop);
-    barrier.arrive_and_wait();
-    segment = (segment + size + direction) % size;
-  }
-}
-
 } // namespace
 
 class RingGroup : public GroupImpl {
@@ -530,50 +540,66 @@ class RingGroup : public GroupImpl {
     // first and accept after.
     if (rank_ < connect_to) {
       log_info(verbose_, "Rank", rank_, "accepting");
-      recv_sockets_ = std::move(accept_connections(nodes[rank_]));
+      sockets_left_ = std::move(accept_connections(nodes[rank_]));
       log_info(verbose_, "Rank", rank_, "connecting to", connect_to);
-      send_sockets_ = std::move(make_connections(nodes[connect_to], verbose));
+      sockets_right_ = std::move(make_connections(nodes[connect_to], verbose));
     } else {
       log_info(verbose_, "Rank", rank_, "connecting to", connect_to);
-      send_sockets_ = std::move(make_connections(nodes[connect_to], verbose));
+      sockets_right_ = std::move(make_connections(nodes[connect_to], verbose));
       log_info(verbose_, "Rank", rank_, "accepting");
-      recv_sockets_ = std::move(accept_connections(nodes[rank_]));
+      sockets_left_ = std::move(accept_connections(nodes[rank_]));
     }
 
-    // Failure if we couldn't make send or recv sockets
-    if (send_sockets_.empty()) {
+    // Failure if we couldn't make right or left sockets
+    if (sockets_right_.empty()) {
       std::ostringstream msg;
-      msg << "[ring] Rank " << rank_ << " has no send sockets.";
+      msg << "[ring] Rank " << rank_ << " has no sockets to the right.";
       throw std::invalid_argument(msg.str());
     }
-    if (recv_sockets_.empty()) {
+    if (sockets_left_.empty()) {
       std::ostringstream msg;
-      msg << "[ring] Rank " << rank_ << " has no recv sockets.";
+      msg << "[ring] Rank " << rank_ << " has no sockets to the left.";
       throw std::invalid_argument(msg.str());
     }
 
     // The following could be relaxed since we can define non-homogeneous rings
     // but it makes things a bit simpler for now.
-    if (send_sockets_.size() != recv_sockets_.size()) {
+    if (sockets_right_.size() != sockets_left_.size()) {
       std::ostringstream msg;
       msg << "[ring] It is required to have as many connections to the left as "
           << "to the right but rank " << rank_ << " has "
-          << send_sockets_.size() << " connections to the right and "
-          << recv_sockets_.size() << " to the left.";
+          << sockets_right_.size() << " connections to the right and "
+          << sockets_left_.size() << " to the left.";
       throw std::invalid_argument(msg.str());
     }
 
-    // Start the necessary threads for completely parallel operation on all
-    // channels. One thread to send, one to receive per socket.
-    pool_.resize(send_sockets_.size() * 2 * 2);
+    // Configure all sockets to use TCP no delay.
+    int one = 1;
+    for (int i = 0; i < sockets_right_.size(); i++) {
+      setsockopt(sockets_right_[i], SOL_TCP, TCP_NODELAY, &one, sizeof(one));
+      setsockopt(sockets_left_[i], SOL_TCP, TCP_NODELAY, &one, sizeof(one));
+    }
+
+    // Start the all reduce threads. One all reduce per direction per ring.
+    pool_.resize(sockets_right_.size() + sockets_left_.size());
+
+    // Create a communication thread per socket. This also converts them to
+    // non-blocking.
+    comm_.add(sockets_right_);
+    comm_.add(sockets_left_);
+
+    // Allocate buffers for the all sum
+    buffers_.resize(
+        (sockets_right_.size() + sockets_left_.size()) * ALL_SUM_BUFFERS *
+        ALL_SUM_SIZE);
   }
 
   ~RingGroup() {
-    for (auto s : send_sockets_) {
+    for (auto s : sockets_right_) {
       shutdown(s, 2);
       close(s);
     }
-    for (auto s : recv_sockets_) {
+    for (auto s : sockets_left_) {
       shutdown(s, 2);
       close(s);
     }
@@ -594,14 +620,45 @@ class RingGroup : public GroupImpl {
   std::shared_ptr<GroupImpl> split(int color, int key = -1) override {
     throw std::runtime_error("[ring] Group split not supported.");
   }
+
   void all_gather(const array& input, array& output) override {
     throw std::runtime_error("[ring] All gather not supported.");
   }
-  void send(const array& input, int dst) override {
-    throw std::runtime_error("[ring] Send not supported.");
+
+  void send(const array& input_, int dst) override {
+    // Make sure that the input is row contiguous
+    array input = ensure_row_contiguous(input_);
+
+    int right = (rank_ + 1) % size_;
+    int left = (rank_ + size_ - 1) % size_;
+    if (dst == right) {
+      send(sockets_right_, input.data<char>(), input.nbytes());
+    } else if (dst == left) {
+      send(sockets_left_, input.data<char>(), input.nbytes());
+    } else {
+      std::ostringstream msg;
+      msg << "[ring] Send only supported to direct neighbors "
+          << "but tried to send to " << dst << " from " << rank_ << std::endl;
+      throw std::runtime_error(msg.str());
+    }
   }
+
   void recv(array& out, int src) override {
-    throw std::runtime_error("[ring] Recv not supported.");
+    // NOTE: We 'll check the sockets with the opposite order of send so that
+    //       they work even with 2 nodes where left and right is the same
+    //       neighbor.
+    int right = (rank_ + 1) % size_;
+    int left = (rank_ + size_ - 1) % size_;
+    if (src == left) {
+      recv(sockets_left_, out.data<char>(), out.nbytes());
+    } else if (src == right) {
+      recv(sockets_right_, out.data<char>(), out.nbytes());
+    } else {
+      std::ostringstream msg;
+      msg << "[ring] Recv only supported from direct neighbors "
+          << "but tried to recv from " << src << " to " << rank_ << std::endl;
+      throw std::runtime_error(msg.str());
+    }
   }
 
  private:
@@ -613,7 +670,8 @@ class RingGroup : public GroupImpl {
     // If the input data cannot be split into size_ segments then copy it and
     // all reduce a local buffer prefilled with 0s.
     if (input.size() < size_) {
-      // TODO: Maybe allocate dynamically so we don't have the constraint below?
+      // TODO: Maybe allocate dynamically so we don't have the constraint
+      // below?
       if (input.itemsize() * size_ > 1024) {
         std::ostringstream msg;
         msg << "Can't perform the ring all reduce of " << output.size()
@@ -621,31 +679,16 @@ class RingGroup : public GroupImpl {
         throw std::runtime_error(msg.str());
       }
 
-      std::future<void> sent, recvd;
-      auto barrier = std::make_unique<Barrier>(2);
       char buffer[1024];
       std::memset(buffer, 0, size_ * input.itemsize());
       std::memcpy(buffer, input.data<char>(), input.nbytes());
-      sent = pool_.enqueue(
-          ring_send<T>,
-          std::reference_wrapper(*barrier),
-          send_sockets_[0],
-          rank_,
-          size_,
-          (T*)buffer,
+      all_sum_impl<T>(
+          reinterpret_cast<T*>(buffers_.data()),
+          reinterpret_cast<T*>(buffer),
           size_,
+          sockets_right_[0],
+          sockets_left_[0],
           -1);
-      recvd = pool_.enqueue(
-          ring_recv_sum<T>,
-          std::reference_wrapper(*barrier),
-          recv_sockets_[0],
-          rank_,
-          size_,
-          (T*)buffer,
-          size_,
-          -1);
-      sent.wait();
-      recvd.wait();
       std::memcpy(output.data<char>(), buffer, output.nbytes());
       return;
     }
@@ -655,137 +698,161 @@ class RingGroup : public GroupImpl {
       std::memcpy(output.data<char>(), input.data<char>(), input.nbytes());
     }
 
-    // All reduce in place. We have `send_channels_.size()` bidirectional
-    // channels so let's split the message up and perform as many parallel
-    // ring-reductions as possible.
-    std::vector<std::future<void>> reductions;
-    std::vector<std::unique_ptr<Barrier>> barriers;
-    size_t packets = ceildiv(output.size(), size_ * PACKET_SIZE);
+    // Split the all reduces so that each member has at least 1 buffer to
+    // send/recv per segment.
+    constexpr size_t min_send_size = 262144;
+    size_t n_reduces = std::max(
+        std::min(
+            sockets_right_.size() + sockets_left_.size(),
+            output.nbytes() / (size_ * min_send_size)),
+        1UL);
+    size_t step = ceildiv(output.size(), n_reduces);
+    std::vector<std::future<void>> all_sums;
 
-    // Large all reduce territory so let's use all we got
-    if (packets >= 2 * send_sockets_.size()) {
-      size_t segment = ceildiv(output.size(), 2 * send_sockets_.size());
-      for (int i = 0; i < send_sockets_.size(); i++) {
-        // 1st ring reduce
-        barriers.emplace_back(std::make_unique<Barrier>(2));
-        reductions.push_back(pool_.enqueue(
-            ring_send<T>,
-            std::reference_wrapper(*barriers.back()),
-            send_sockets_[i],
-            rank_,
-            size_,
-            output.data<T>() + 2 * i * segment,
-            std::min(output.size() - 2 * i * segment, segment),
-            -1));
-        reductions.push_back(pool_.enqueue(
-            ring_recv_sum<T>,
-            std::reference_wrapper(*barriers.back()),
-            recv_sockets_[i],
-            rank_,
-            size_,
-            output.data<T>() + 2 * i * segment,
-            std::min(output.size() - 2 * i * segment, segment),
-            -1));
+    for (int i = 0; i < n_reduces; i++) {
+      all_sums.emplace_back(pool_.enqueue(std::bind(
+          &RingGroup::all_sum_impl<T>,
+          this,
+          reinterpret_cast<T*>(
+              buffers_.data() + i * ALL_SUM_SIZE * ALL_SUM_BUFFERS),
+          output.data<T>() + i * step,
+          std::min(output.size(), (i + 1) * step) - i * step,
+          sockets_right_[i / 2],
+          sockets_left_[i / 2],
+          (i % 2) ? -1 : 1)));
+    }
+    for (auto& f : all_sums) {
+      f.wait();
+    }
+  }
 
-        // 2nd ring reduce
-        barriers.emplace_back(std::make_unique<Barrier>(2));
-        reductions.push_back(pool_.enqueue(
-            ring_send<T>,
-            std::reference_wrapper(*barriers.back()),
-            recv_sockets_[i],
-            rank_,
-            size_,
-            output.data<T>() + (2 * i + 1) * segment,
-            std::min(output.size() - (2 * i + 1) * segment, segment),
-            1));
-        reductions.push_back(pool_.enqueue(
-            ring_recv_sum<T>,
-            std::reference_wrapper(*barriers.back()),
-            send_sockets_[i],
-            rank_,
-            size_,
-            output.data<T>() + (2 * i + 1) * segment,
-            std::min(output.size() - (2 * i + 1) * segment, segment),
-            1));
+  template <typename T>
+  void all_sum_impl(
+      T* buffer,
+      T* data,
+      size_t data_size,
+      int socket_right,
+      int socket_left,
+      int direction) {
+    // Choose which socket we send to and recv from
+    int socket_send = (direction < 0) ? socket_right : socket_left;
+    int socket_recv = (direction < 0) ? socket_left : socket_right;
+
+    // We split the data into `size_` segments of size `segment_size` and each
+    // of these in smaller segments of ALL_SUM_SIZE which we 'll call packets.
+    size_t segment_size = ceildiv(data_size, size_);
+    size_t BUFFER_SIZE =
+        std::max(32768UL, std::min(ALL_SUM_SIZE / sizeof(T), segment_size / 2));
+    size_t n_packets = ceildiv(segment_size, BUFFER_SIZE);
+
+    // Initial segments
+    int send_segment = rank_;
+    int recv_segment = (rank_ + direction + size_) % size_;
+
+    // Plan the whole reduce in terms of sends and recvs as indices in data.
+    // It makes the actual async send and recv a bit simpler to follow when
+    // there are less offset calculations around.
+    std::vector<std::pair<size_t, size_t>> send_plan;
+    std::vector<std::pair<size_t, size_t>> recv_plan;
+
+    // Two times the same send/recv operations, first scatter reduce and then
+    // gather.
+    for (int k = 0; k < 2; k++) {
+      for (int i = 0; i < size_ - 1; i++) {
+        size_t send_start = send_segment * segment_size;
+        size_t send_stop =
+            std::min((send_segment + 1) * segment_size, data_size);
+        size_t recv_start = recv_segment * segment_size;
+        size_t recv_stop =
+            std::min((recv_segment + 1) * segment_size, data_size);
+
+        for (size_t j = 0; j < n_packets; j++) {
+          send_plan.emplace_back(
+              std::min(send_start + j * BUFFER_SIZE, send_stop),
+              std::min(send_start + (j + 1) * BUFFER_SIZE, send_stop));
+          recv_plan.emplace_back(
+              std::min(recv_start + j * BUFFER_SIZE, recv_stop),
+              std::min(recv_start + (j + 1) * BUFFER_SIZE, recv_stop));
+        }
+
+        send_segment = (send_segment + size_ + direction) % size_;
+        recv_segment = (recv_segment + size_ + direction) % size_;
       }
     }
 
-    // At least 2 reductions so we can be from small to medium
-    else if (packets > 1) {
-      size_t segment = ceildiv(output.size(), packets);
-      for (int i = 0; i < send_sockets_.size(); i++) {
-        barriers.emplace_back(std::make_unique<Barrier>(2));
-        reductions.push_back(pool_.enqueue(
-            ring_send<T>,
-            std::reference_wrapper(*barriers.back()),
-            send_sockets_[i],
-            rank_,
-            size_,
-            output.data<T>() + i * segment,
-            std::min(output.size() - i * segment, segment),
-            -1));
-        reductions.push_back(pool_.enqueue(
-            ring_recv_sum<T>,
-            std::reference_wrapper(*barriers.back()),
-            recv_sockets_[i],
-            rank_,
-            size_,
-            output.data<T>() + i * segment,
-            std::min(output.size() - i * segment, segment),
-            -1));
-      }
-      for (int i = 0; i < packets - send_sockets_.size(); i++) {
-        barriers.emplace_back(std::make_unique<Barrier>(2));
-        reductions.push_back(pool_.enqueue(
-            ring_send<T>,
-            std::reference_wrapper(*barriers.back()),
-            recv_sockets_[i],
-            rank_,
-            size_,
-            output.data<T>() + (send_sockets_.size() + i) * segment,
-            std::min(
-                output.size() - (send_sockets_.size() + i) * segment, segment),
-            1));
-        reductions.push_back(pool_.enqueue(
-            ring_recv_sum<T>,
-            std::reference_wrapper(*barriers.back()),
-            send_sockets_[i],
-            rank_,
-            size_,
-            output.data<T>() + (send_sockets_.size() + i) * segment,
-            std::min(
-                output.size() - (send_sockets_.size() + i) * segment, segment),
-            1));
-      }
+    // Running the plan is fairly simple, we keep a send and a recv in flight
+    // while doing the summation.
+    T* recv_buffers[ALL_SUM_BUFFERS];
+    for (int i = 0; i < ALL_SUM_BUFFERS; i++) {
+      recv_buffers[i] = buffer + i * BUFFER_SIZE;
     }
+    std::future<void> sends[2], recvs[2];
+    int a = 0;
+    int b = (n_packets > 1) ? 1 : 0;
+    for (int i = 0, j = -b; i < send_plan.size(); j++, i++) {
+      sends[a] = comm_.send(
+          socket_send,
+          data + send_plan[i].first,
+          send_plan[i].second - send_plan[i].first);
+      if (2 * i < send_plan.size()) {
+        recvs[a] = comm_.recv(
+            socket_recv,
+            recv_buffers[i % ALL_SUM_BUFFERS],
+            recv_plan[i].second - recv_plan[i].first);
+      } else {
+        recvs[a] = comm_.recv(
+            socket_recv,
+            data + recv_plan[i].first,
+            recv_plan[i].second - recv_plan[i].first);
+      }
 
-    // Small reduction which won't really benefit much from parallelization.
-    // TODO: Verify that this is true cause PACKET_SIZE * size_ can still be a
-    //       fairly large array.
-    else {
-      barriers.emplace_back(std::make_unique<Barrier>(2));
-      reductions.push_back(pool_.enqueue(
-          ring_send<T>,
-          std::reference_wrapper(*barriers.back()),
-          send_sockets_[0],
-          rank_,
-          size_,
-          output.data<T>(),
-          output.size(),
-          -1));
-      reductions.push_back(pool_.enqueue(
-          ring_recv_sum<T>,
-          std::reference_wrapper(*barriers.back()),
-          recv_sockets_[0],
-          rank_,
-          size_,
-          output.data<T>(),
-          output.size(),
-          -1));
+      if (j >= 0) {
+        sends[b].wait();
+        recvs[b].wait();
+        if (2 * j < send_plan.size()) {
+          sum_inplace<T>(
+              recv_buffers[j % ALL_SUM_BUFFERS],
+              data + recv_plan[j].first,
+              recv_plan[j].second - recv_plan[j].first);
+        }
+      }
+
+      std::swap(a, b);
     }
+    sends[b].wait();
+    recvs[b].wait();
+  }
 
-    // Wait for the reductions to finish.
-    for (auto& f : reductions) {
+  void send(const std::vector<int>& sockets, char* data, size_t data_size) {
+    size_t segment_size = std::max(1024UL, ceildiv(data_size, sockets.size()));
+    std::vector<std::future<void>> sends;
+    for (int i = 0; i < sockets.size(); i++) {
+      if (i * segment_size >= data_size) {
+        break;
+      }
+      sends.emplace_back(comm_.send(
+          sockets[i],
+          data + i * segment_size,
+          std::min(data_size, (i + 1) * segment_size) - i * segment_size));
+    }
+    for (auto& f : sends) {
+      f.wait();
+    }
+  }
+
+  void recv(const std::vector<int>& sockets, char* data, size_t data_size) {
+    size_t segment_size = std::max(1024UL, ceildiv(data_size, sockets.size()));
+    std::vector<std::future<void>> recvs;
+    for (int i = 0; i < sockets.size(); i++) {
+      if (i * segment_size >= data_size) {
+        break;
+      }
+      recvs.emplace_back(comm_.recv(
+          sockets[i],
+          data + i * segment_size,
+          std::min(data_size, (i + 1) * segment_size) - i * segment_size));
+    }
+    for (auto& f : recvs) {
       f.wait();
     }
   }
@@ -796,9 +863,12 @@ class RingGroup : public GroupImpl {
   bool verbose_;
 
   ThreadPool pool_;
+  CommunicationThreads comm_;
 
-  std::vector<int> send_sockets_;
-  std::vector<int> recv_sockets_;
+  std::vector<int> sockets_right_;
+  std::vector<int> sockets_left_;
+
+  std::vector<char> buffers_;
 };
 
 bool is_available() {

mlx/dtype.cpp

@@ -44,7 +44,7 @@ constexpr Dtype type_rules[num_types][num_types] = {
   {int64,     int64,     int64,     int64,     float32,   int64,     int64,     int64,     int64,     float16,   float32,  float64,   bfloat16,  complex64}, // int64
   {float16,   float16,   float16,   float16,   float16,   float16,   float16,   float16,   float16,   float16,   float32,  float64,   float32,   complex64}, // float16
   {float32,   float32,   float32,   float32,   float32,   float32,   float32,   float32,   float32,   float32,   float32,  float64,   float32,   complex64}, // float32
-  {float64,   float32,   float32,   float32,   float32,   float32,   float32,   float32,   float32,   float32,   float32,  float64,   float32,   complex64}, // float64
+  {float64,   float64,   float64,   float64,   float64,   float64,   float64,   float64,   float64,   float64,   float64,  float64,   float64,   complex64}, // float64
   {bfloat16,  bfloat16,  bfloat16,  bfloat16,  bfloat16,  bfloat16,  bfloat16,  bfloat16,  bfloat16,  float32,   float32,  float64,   bfloat16,  complex64}, // bfloat16
   {complex64, complex64, complex64, complex64, complex64, complex64, complex64, complex64, complex64, complex64, complex64,complex64, complex64, complex64}, // complex64
 };

mlx/export.cpp

@@ -4,9 +4,7 @@
 #include "mlx/fast_primitives.h"
 #include "mlx/primitives.h"
 #include "mlx/utils.h"
-
-#define STRINGIFY(x) #x
-#define TOSTRING(x) STRINGIFY(x)
+#include "mlx/version.h"
 
 // clang-format off
 #define SERIALIZE_PRIMITIVE(primitive, ...)  \
@@ -379,7 +377,7 @@ struct PrimitiveFactory {
 };
 
 void write_header(Writer& os, int count, bool shapeless) {
-  serialize(os, std::string(TOSTRING(MLX_VERSION)));
+  serialize(os, std::string(version()));
   serialize(os, count);
   serialize(os, shapeless);
 }

mlx/fast.cpp

@@ -54,30 +54,34 @@ std::pair<std::vector<array>, std::vector<int>> Custom::vmap(
 
 array rms_norm(
     const array& x,
-    const array& weight,
+    const std::optional<array>& weight,
     float eps,
     StreamOrDevice s_ /* = {} */) {
+  bool has_weight = weight.has_value();
+
   if (x.ndim() == 0) {
     std::ostringstream msg;
     msg << "[rms_norm] Input must have at least 1 dimension but got input with "
            "0 dimensions.";
     throw std::invalid_argument(msg.str());
   }
-  if (weight.ndim() != 1) {
-    std::ostringstream msg;
-    msg << "[rms_norm] weight must have 1 dimension but has " << weight.ndim()
-        << " dimensions.";
-    throw std::invalid_argument(msg.str());
-  }
-  if (weight.size() != x.shape(-1)) {
-    std::ostringstream msg;
-    msg << "[rms_norm] weight must have the same size as the last dimension of"
-           " x but has "
-        << weight.size() << " elements.";
-    throw std::invalid_argument(msg.str());
+  if (has_weight) {
+    if ((*weight).ndim() != 1) {
+      std::ostringstream msg;
+      msg << "[rms_norm] (*weight) must have 1 dimension but has "
+          << (*weight).ndim() << " dimensions.";
+      throw std::invalid_argument(msg.str());
+    }
+    if ((*weight).size() != x.shape(-1)) {
+      std::ostringstream msg;
+      msg << "[rms_norm] (*weight) must have the same size as the last dimension of"
+             " x but has "
+          << (*weight).size() << " elements.";
+      throw std::invalid_argument(msg.str());
+    }
   }
 
-  auto out_type = result_type(x, weight);
+  auto out_type = (weight.has_value()) ? result_type(x, (*weight)) : x.dtype();
   if (!issubdtype(out_type, floating)) {
     std::ostringstream msg;
     msg << "[rms_norm] Received unsupported type " << out_type << ".";
@@ -85,27 +89,36 @@ array rms_norm(
   }
 
   auto s = to_stream(s_);
-  auto fallback = [eps, out_type, s](const std::vector<array>& inputs) {
-    auto x = astype(inputs[0], float32, s);
-    x = multiply(
-        x,
-        rsqrt(
-            add(mean(square(x, s), -1, /* keepdims */ true, s),
-                array(eps, float32),
+  auto fallback =
+      [has_weight, eps, out_type, s](const std::vector<array>& inputs) {
+        auto x = astype(inputs[0], float32, s);
+        x = multiply(
+            x,
+            rsqrt(
+                add(mean(square(x, s), -1, /* keepdims */ true, s),
+                    array(eps, float32),
+                    s),
                 s),
-            s),
-        s);
-    x = astype(x, out_type, s);
-    return std::vector<array>{multiply(inputs[1], x, s)};
-  };
+            s);
+        x = astype(x, out_type, s);
+
+        if (has_weight) {
+          x = multiply(x, inputs[1], s);
+        }
+
+        return std::vector<array>{x};
+      };
+
+  auto passed_weight =
+      (has_weight) ? astype(*weight, out_type, s) : array(1, out_type);
   if (s.device == Device::gpu) {
     return array(
         x.shape(),
         out_type,
         std::make_shared<RMSNorm>(s, fallback, eps),
-        {astype(x, out_type, s), astype(weight, out_type, s)});
+        {astype(x, out_type, s), passed_weight});
   }
-  return fallback({x, weight})[0];
+  return fallback({x, passed_weight})[0];
 }
 
 std::vector<array> RMSNorm::vjp(
@@ -141,8 +154,12 @@ std::vector<array> RMSNorm::vjp(
     // df/dw
     std::vector<int> axes(g.ndim() - 1);
     std::iota(axes.begin(), axes.end(), 0);
-    vjps.push_back(
-        sum(multiply(g, multiply(x, n, s), s), axes, /* keepdims= */ false, s));
+    if (w.ndim() == 0) {
+      vjps.push_back(zeros_like(w, s));
+    } else {
+      vjps.push_back(sum(
+          multiply(g, multiply(x, n, s), s), axes, /* keepdims= */ false, s));
+    }
 
     return vjps;
   };
@@ -177,28 +194,30 @@ array layer_norm(
     const std::optional<array>& bias,
     float eps,
     StreamOrDevice s_ /* = {} */) {
+  bool has_weight = weight.has_value();
+  bool has_bias = bias.has_value();
+
   if (x.ndim() == 0) {
     std::ostringstream msg;
     msg << "[layer_norm] Input must have at least 1 dimension but got input with "
            "0 dimensions.";
     throw std::invalid_argument(msg.str());
   }
-  if (weight.has_value() && (*weight).ndim() != 1) {
+  if (has_weight && (*weight).ndim() != 1) {
     std::ostringstream msg;
     msg << "[layer_norm] weight must have 1 dimension but has "
         << (*weight).ndim() << " dimensions.";
     throw std::invalid_argument(msg.str());
   }
-  if (bias.has_value() && (*bias).ndim() != 1) {
+  if (has_bias && (*bias).ndim() != 1) {
     std::ostringstream msg;
     msg << "[layer_norm] bias must have 1 dimension but has " << (*bias).ndim()
         << " dimensions.";
     throw std::invalid_argument(msg.str());
   }
 
-  auto out_type = (weight.has_value())
-      ? ((bias.has_value()) ? result_type(x, *weight, *bias)
-                            : result_type(x, *weight))
+  auto out_type = (has_weight)
+      ? ((has_bias) ? result_type(x, *weight, *bias) : result_type(x, *weight))
       : x.dtype();
   if (!issubdtype(out_type, floating)) {
     std::ostringstream msg;
@@ -207,8 +226,6 @@ array layer_norm(
   }
 
   auto s = to_stream(s_);
-  bool has_weight = weight.has_value();
-  bool has_bias = bias.has_value();
   auto fallback = [has_weight, has_bias, eps, out_type, s](
                       const std::vector<array>& inputs) {
     auto x = astype(inputs[0], float32, s);
@@ -234,9 +251,9 @@ array layer_norm(
   };
 
   auto passed_weight =
-      astype((weight.has_value()) ? *weight : array(1, out_type), out_type);
+      (has_weight) ? astype(*weight, out_type, s) : array(1, out_type);
   auto passed_bias =
-      astype((bias.has_value()) ? *bias : array(0, out_type), out_type);
+      (has_bias) ? astype(*bias, out_type, s) : array(0, out_type);
 
   if (s.device == Device::gpu) {
     return array(
@@ -348,6 +365,11 @@ array rope(
         << x.ndim() << " dimensions.";
     throw std::invalid_argument(msg.str());
   }
+  if (!issubdtype(x.dtype(), floating)) {
+    std::ostringstream msg;
+    msg << "[rope] Input must be a floating type but got " << x.dtype() << ".";
+    throw std::invalid_argument(msg.str());
+  }
   if (offset.size() != 1) {
     std::ostringstream msg;
     msg << "[rope] offset must be a scalar but has shape " << offset.shape()
@@ -688,12 +710,13 @@ array scaled_dot_product_attention(
       query_head_dim == value_head_dim &&
       (query_head_dim == 64 || query_head_dim == 96 || query_head_dim == 128);
   const bool sdpa_full_supported_head_dim = query_head_dim == value_head_dim &&
-      (query_head_dim == 64 || query_head_dim == 80);
+      (query_head_dim == 64 || query_head_dim == 80 || query_head_dim == 128);
 
   const bool supports_sdpa_full = query_sequence_length >= threshold && !mask &&
       sdpa_full_supported_head_dim && stream.device == Device::gpu;
 
-  const bool supports_sdpa_vector = query_sequence_length == 1 &&
+  const bool supports_sdpa_vector = (query_sequence_length <= 8) &&
+      (query_sequence_length <= k.shape(-2)) &&
       (!mask || mask->dtype() == bool_) && sdpa_vector_supported_head_dim &&
       stream.device == Device::gpu;
 
@@ -804,14 +827,17 @@ affine_quantize(const array& w, int group_size, int bits, StreamOrDevice s_) {
     auto wshape = w.shape();
     wshape.back() = -1;
 
-    array zero(0, w.dtype());
-    array n_bins((1 << bits) - 1, w.dtype()); // 2**bits - 1
-    array eps(1e-7, w.dtype());
+    array zero(0, float32);
+    array n_bins((1 << bits) - 1, float32); // 2**bits - 1
+    array eps(1e-7, float32);
 
     array packed_w = reshape(w, {-1, 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);
+    w_max = astype(w_max, float32, s);
+    w_min = astype(w_min, float32, s);
+
     array mask = greater(abs(w_min, s), abs(w_max, s), s);
     array scales =
         maximum(divide(subtract(w_max, w_min, s), n_bins, s), eps, s);
@@ -822,6 +848,9 @@ affine_quantize(const array& w, int group_size, int bits, StreamOrDevice s_) {
     array biases = where(equal(q0, zero, s), zero, edge, s);
 
     packed_w = pack_and_quantize(packed_w, scales, biases, bits, s);
+
+    scales = astype(scales, w.dtype(), s);
+    biases = astype(biases, w.dtype(), s);
     return {
         reshape(packed_w, wshape, s),
         reshape(scales, wshape, s),

mlx/fast.h

@@ -10,7 +10,7 @@ namespace mlx::core::fast {
 
 array rms_norm(
     const array& x,
-    const array& weight,
+    const std::optional<array>& weight,
     float eps,
     StreamOrDevice s = {});
 

mlx/linalg.cpp

@@ -18,6 +18,14 @@ void check_cpu_stream(const StreamOrDevice& s, const std::string& prefix) {
         "Explicitly pass a CPU stream to run it.");
   }
 }
+void check_float(Dtype dtype, const std::string& prefix) {
+  if (dtype != float32 && dtype != float64) {
+    std::ostringstream msg;
+    msg << prefix << " Arrays must have type float32 or float64. "
+        << "Received array with type " << dtype << ".";
+    throw std::invalid_argument(msg.str());
+  }
+}
 
 Dtype at_least_float(const Dtype& d) {
   return issubdtype(d, inexact) ? d : promote_types(d, float32);
@@ -94,8 +102,21 @@ inline array matrix_norm(
         dtype,
         s);
   } else if (ord == 2.0 || ord == -2.0) {
-    throw std::runtime_error(
-        "[linalg::norm] Singular value norms are not implemented.");
+    row_axis = (axis[0] < 0) ? axis[0] + a.ndim() : axis[0];
+    col_axis = (axis[1] < 0) ? axis[1] + a.ndim() : axis[1];
+    auto a_matrix = (row_axis > col_axis)
+        ? moveaxis(moveaxis(a, row_axis, -1, s), col_axis, -1, s)
+        : moveaxis(moveaxis(a, col_axis, -1, s), row_axis, -2, s);
+    a_matrix = svd(a_matrix, false, s).at(0);
+    a_matrix = (ord == 2.0) ? max(a_matrix, -1, false, s)
+                            : min(a_matrix, -1, false, s);
+    if (keepdims) {
+      std::vector<int> sorted_axes = (row_axis < col_axis)
+          ? std::vector<int>{row_axis, col_axis}
+          : std::vector<int>{col_axis, row_axis};
+      a_matrix = expand_dims(a_matrix, sorted_axes, s);
+    }
+    return astype(a_matrix, dtype, s);
   } else {
     std::ostringstream msg;
     msg << "[linalg::norm] Invalid ord " << ord << " for matrix norm.";
@@ -112,8 +133,19 @@ inline array matrix_norm(
   if (ord == "f" || ord == "fro") {
     return l2_norm(a, axis, keepdims, s);
   } else if (ord == "nuc") {
-    throw std::runtime_error(
-        "[linalg::norm] Nuclear norm not yet implemented.");
+    int row_axis = (axis[0] < 0) ? axis[0] + a.ndim() : axis[0];
+    int col_axis = (axis[1] < 0) ? axis[1] + a.ndim() : axis[1];
+    auto a_matrix = (row_axis > col_axis)
+        ? moveaxis(moveaxis(a, row_axis, -1, s), col_axis, -1, s)
+        : moveaxis(moveaxis(a, col_axis, -1, s), row_axis, -2, s);
+    a_matrix = sum(svd(a_matrix, false, s).at(0), -1, false, s);
+    if (keepdims) {
+      std::vector<int> sorted_axes = (row_axis < col_axis)
+          ? std::vector<int>{row_axis, col_axis}
+          : std::vector<int>{col_axis, row_axis};
+      a_matrix = expand_dims(a_matrix, sorted_axes, s);
+    }
+    return a_matrix;
   } else {
     std::ostringstream msg;
     msg << "[linalg::norm] Invalid ord value '" << ord << "' for matrix norm.";
@@ -184,12 +216,8 @@ array norm(
 
 std::pair<array, array> qr(const array& a, StreamOrDevice s /* = {} */) {
   check_cpu_stream(s, "[linalg::qr]");
-  if (a.dtype() != float32) {
-    std::ostringstream msg;
-    msg << "[linalg::qr] Arrays must type float32. Received array "
-        << "with type " << a.dtype() << ".";
-    throw std::invalid_argument(msg.str());
-  }
+  check_float(a.dtype(), "[linalg::qr]");
+
   if (a.ndim() < 2) {
     std::ostringstream msg;
     msg << "[linalg::qr] Arrays must have >= 2 dimensions. Received array "
@@ -210,14 +238,11 @@ std::pair<array, array> qr(const array& a, StreamOrDevice s /* = {} */) {
   return std::make_pair(out[0], out[1]);
 }
 
-std::vector<array> svd(const array& a, StreamOrDevice s /* = {} */) {
+std::vector<array>
+svd(const array& a, bool compute_uv, StreamOrDevice s /* = {} */) {
   check_cpu_stream(s, "[linalg::svd]");
-  if (a.dtype() != float32) {
-    std::ostringstream msg;
-    msg << "[linalg::svd] Input array must have type float32. Received array "
-        << "with type " << a.dtype() << ".";
-    throw std::invalid_argument(msg.str());
-  }
+  check_float(a.dtype(), "[linalg::svd]");
+
   if (a.ndim() < 2) {
     std::ostringstream msg;
     msg << "[linalg::svd] Input array must have >= 2 dimensions. Received array "
@@ -230,14 +255,22 @@ std::vector<array> svd(const array& a, StreamOrDevice s /* = {} */) {
   const auto n = a.shape(-1);
   const auto rank = a.ndim();
 
-  auto u_shape = a.shape();
-  u_shape[rank - 2] = m;
-  u_shape[rank - 1] = m;
-
   auto s_shape = a.shape();
   s_shape.pop_back();
   s_shape[rank - 2] = std::min(m, n);
 
+  if (!compute_uv) {
+    return {array(
+        std::move(s_shape),
+        std::move(a.dtype()),
+        std::make_shared<SVD>(to_stream(s), compute_uv),
+        {a})};
+  }
+
+  auto u_shape = a.shape();
+  u_shape[rank - 2] = m;
+  u_shape[rank - 1] = m;
+
   auto vt_shape = a.shape();
   vt_shape[rank - 2] = n;
   vt_shape[rank - 1] = n;
@@ -245,18 +278,14 @@ std::vector<array> svd(const array& a, StreamOrDevice s /* = {} */) {
   return array::make_arrays(
       {u_shape, s_shape, vt_shape},
       {a.dtype(), a.dtype(), a.dtype()},
-      std::make_shared<SVD>(to_stream(s)),
+      std::make_shared<SVD>(to_stream(s), compute_uv),
       {a});
 }
 
 array inv_impl(const array& a, bool tri, bool upper, StreamOrDevice s) {
   check_cpu_stream(s, "[linalg::inv]");
-  if (a.dtype() != float32) {
-    std::ostringstream msg;
-    msg << "[linalg::inv] Arrays must type float32. Received array "
-        << "with type " << a.dtype() << ".";
-    throw std::invalid_argument(msg.str());
-  }
+  check_float(a.dtype(), "[linalg::inv]");
+
   if (a.ndim() < 2) {
     std::ostringstream msg;
     msg << "[linalg::inv] Arrays must have >= 2 dimensions. Received array "
@@ -292,13 +321,7 @@ array cholesky(
     bool upper /* = false */,
     StreamOrDevice s /* = {} */) {
   check_cpu_stream(s, "[linalg::cholesky]");
-  if (a.dtype() != float32) {
-    std::ostringstream msg;
-    msg << "[linalg::cholesky] Arrays must type float32. Received array "
-        << "with type " << a.dtype() << ".";
-    throw std::invalid_argument(msg.str());
-  }
-
+  check_float(a.dtype(), "[linalg::cholesky]");
   if (a.ndim() < 2) {
     std::ostringstream msg;
     msg << "[linalg::cholesky] Arrays must have >= 2 dimensions. Received array "
@@ -321,12 +344,8 @@ array cholesky(
 
 array pinv(const array& a, StreamOrDevice s /* = {} */) {
   check_cpu_stream(s, "[linalg::pinv]");
-  if (a.dtype() != float32) {
-    std::ostringstream msg;
-    msg << "[linalg::pinv] Arrays must type float32. Received array "
-        << "with type " << a.dtype() << ".";
-    throw std::invalid_argument(msg.str());
-  }
+  check_float(a.dtype(), "[linalg::pinv]");
+
   if (a.ndim() < 2) {
     std::ostringstream msg;
     msg << "[linalg::pinv] Arrays must have >= 2 dimensions. Received array "
@@ -337,7 +356,7 @@ array pinv(const array& a, StreamOrDevice s /* = {} */) {
   int m = a.shape(-2);
   int n = a.shape(-1);
   int k = std::min(m, n);
-  auto outs = linalg::svd(a, s);
+  auto outs = linalg::svd(a, true, s);
   array U = outs[0];
   array S = outs[1];
   array V = outs[2];
@@ -368,12 +387,7 @@ array cholesky_inv(
     bool upper /* = false */,
     StreamOrDevice s /* = {} */) {
   check_cpu_stream(s, "[linalg::cholesky_inv]");
-  if (L.dtype() != float32) {
-    std::ostringstream msg;
-    msg << "[linalg::cholesky_inv] Arrays must type float32. Received array "
-        << "with type " << L.dtype() << ".";
-    throw std::invalid_argument(msg.str());
-  }
+  check_float(L.dtype(), "[linalg::cholesky_inv]");
 
   if (L.ndim() < 2) {
     std::ostringstream msg;
@@ -474,12 +488,7 @@ void validate_eigh(
     const StreamOrDevice& stream,
     const std::string fname) {
   check_cpu_stream(stream, fname);
-  if (a.dtype() != float32) {
-    std::ostringstream msg;
-    msg << fname << " Arrays must have type float32. Received array "
-        << "with type " << a.dtype() << ".";
-    throw std::invalid_argument(msg.str());
-  }
+  check_float(a.dtype(), fname);
 
   if (a.ndim() < 2) {
     std::ostringstream msg;
@@ -524,12 +533,7 @@ void validate_lu(
     const StreamOrDevice& stream,
     const std::string& fname) {
   check_cpu_stream(stream, fname);
-  if (a.dtype() != float32) {
-    std::ostringstream msg;
-    msg << fname << " Arrays must type float32. Received array "
-        << "with type " << a.dtype() << ".";
-    throw std::invalid_argument(msg.str());
-  }
+  check_float(a.dtype(), fname);
 
   if (a.ndim() < 2) {
     std::ostringstream msg;
@@ -539,10 +543,6 @@ void validate_lu(
         << a.ndim() << " dimensions.";
     throw std::invalid_argument(msg.str());
   }
-
-  if (a.shape(-1) != a.shape(-2)) {
-    throw std::invalid_argument(fname + " Only defined for square matrices.");
-  }
 }
 
 std::vector<array> lu_helper(const array& a, StreamOrDevice s /* = {} */) {
@@ -552,8 +552,10 @@ std::vector<array> lu_helper(const array& a, StreamOrDevice s /* = {} */) {
   Shape pivots_shape(a.shape().begin(), a.shape().end() - 2);
   pivots_shape.push_back(std::min(m, n));
 
+  Shape row_idx_shape(a.shape().begin(), a.shape().end() - 1);
+
   return array::make_arrays(
-      {a.shape(), pivots_shape, pivots_shape},
+      {a.shape(), pivots_shape, row_idx_shape},
       {a.dtype(), uint32, uint32},
       std::make_shared<LUF>(to_stream(s)),
       {astype(a, a.dtype(), s)});
@@ -565,10 +567,24 @@ std::vector<array> lu(const array& a, StreamOrDevice s /* = {} */) {
   auto out = lu_helper(a, s);
   auto& LU = out[0];
   auto& row_pivots = out[2];
-
-  int N = a.shape(-1);
-  auto L = add(tril(LU, /* k = */ -1, s), eye(N, s), s);
+  auto L = tril(LU, /* k = */ -1, s);
   auto U = triu(LU, /* k = */ 0, s);
+
+  int M = a.shape(-2);
+  int N = a.shape(-1);
+  int K = std::min(M, N);
+  if (N != K) {
+    auto start = Shape(L.ndim(), 0);
+    auto stop = L.shape();
+    stop.back() = K;
+    L = slice(L, std::move(start), std::move(stop), s);
+  } else if (M != K) {
+    auto start = Shape(U.ndim(), 0);
+    auto stop = U.shape();
+    stop[U.ndim() - 2] = K;
+    U = slice(U, std::move(start), std::move(stop), s);
+  }
+  L = add(L, eye(M, K, s), s);
   return {row_pivots, L, U};
 }
 
@@ -615,10 +631,12 @@ void validate_solve(
   }
 
   auto out_type = promote_types(a.dtype(), b.dtype());
-  if (out_type != float32) {
+  if (out_type != float32 && out_type != float64) {
     std::ostringstream msg;
-    msg << fname << " Input arrays must promote to float32. Received arrays "
-        << "with type " << a.dtype() << " and " << b.dtype() << ".";
+    msg << fname
+        << " Input arrays must promote to float32 or float64. "
+           " Received arrays with type "
+        << a.dtype() << " and " << b.dtype() << ".";
     throw std::invalid_argument(msg.str());
   }
 }

mlx/linalg.h

@@ -62,7 +62,11 @@ norm(const array& a, int axis, bool keepdims = false, StreamOrDevice s = {}) {
 
 std::pair<array, array> qr(const array& a, StreamOrDevice s = {});
 
-std::vector<array> svd(const array& a, StreamOrDevice s = {});
+std::vector<array>
+svd(const array& a, bool compute_uv, StreamOrDevice s /* = {} */);
+inline std::vector<array> svd(const array& a, StreamOrDevice s = {}) {
+  return svd(a, true, s);
+}
 
 array inv(const array& a, StreamOrDevice s = {});
 

mlx/mlx.h

@@ -19,3 +19,4 @@
 #include "mlx/stream.h"
 #include "mlx/transforms.h"
 #include "mlx/utils.h"
+#include "mlx/version.h"

mlx/ops.cpp

@@ -230,16 +230,16 @@ array linspace(
     throw std::invalid_argument(msg.str());
   }
   if (num == 1) {
-    return astype(array({start}), dtype, to_stream(s));
+    return astype(array({start}), dtype, s);
   }
-  array sequence = arange(0, num, float32, to_stream(s));
-  float step = (stop - start) / (num - 1);
+  array t = divide(arange(0, num, float32, s), array(num - 1, float32), s);
+  array t_bar = subtract(array(1, float32), t, s);
   return astype(
-      add(multiply(sequence, array(step), to_stream(s)),
-          array(start),
-          to_stream(s)),
+      add(multiply(t_bar, array(start, float32), s),
+          multiply(t, array(stop, float32), s),
+          s),
       dtype,
-      to_stream(s));
+      s);
 }
 
 array astype(array a, Dtype dtype, StreamOrDevice s /* = {} */) {
@@ -3882,7 +3882,7 @@ array conv_general(
 
   return array(
       std::move(out_shape),
-      out_type,
+      in.dtype(),
       std::make_shared<Convolution>(
           to_stream(s),
           stride,

mlx/primitives.cpp

@@ -4940,7 +4940,8 @@ std::pair<std::vector<array>, std::vector<int>> SVD::vmap(
     const std::vector<int>& axes) {
   auto ax = axes[0] >= 0 ? 0 : -1;
   auto a = axes[0] > 0 ? moveaxis(inputs[0], axes[0], 0, stream()) : inputs[0];
-  return {{linalg::svd(a, stream())}, {ax, ax, ax}};
+  std::vector<int> new_axes(compute_uv_ ? 3 : 1, ax);
+  return {linalg::svd(a, compute_uv_, stream()), std::move(new_axes)};
 }
 
 std::pair<std::vector<array>, std::vector<int>> Inverse::vmap(

mlx/primitives.h

@@ -2287,7 +2287,8 @@ class QRF : public Primitive {
 /* SVD primitive. */
 class SVD : public Primitive {
  public:
-  explicit SVD(Stream stream) : Primitive(stream) {}
+  explicit SVD(Stream stream, bool compute_uv)
+      : Primitive(stream), compute_uv_(compute_uv) {}
 
   void eval_cpu(const std::vector<array>& inputs, std::vector<array>& outputs)
       override;
@@ -2296,6 +2297,12 @@ class SVD : public Primitive {
 
   DEFINE_VMAP()
   DEFINE_PRINT(SVD)
+  auto state() const {
+    return compute_uv_;
+  }
+
+ private:
+  bool compute_uv_;
 };
 
 /* Matrix inversion primitive. */

mlx/random.cpp

@@ -244,7 +244,7 @@ array multivariate_normal(
 
   // Compute the square-root of the covariance matrix, using the SVD
   auto covariance = astype(cov, float32, stream);
-  auto SVD = linalg::svd(covariance, stream);
+  auto SVD = linalg::svd(covariance, true, stream);
   auto std = astype(
       matmul(
           multiply(

mlx/version.cpp

@@ -0,0 +1,16 @@
+// Copyright © 2025 Apple Inc.
+
+#include <string>
+
+#include "mlx/version.h"
+
+#define STRINGIFY(x) #x
+#define TOSTRING(x) STRINGIFY(x)
+
+namespace mlx::core {
+
+std::string version() {
+  return TOSTRING(MLX_VERSION);
+}
+
+} // namespace mlx::core

mlx/version.h

@@ -0,0 +1,20 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#define MLX_VERSION_MAJOR 0
+#define MLX_VERSION_MINOR 23
+#define MLX_VERSION_PATCH 2
+#define MLX_VERSION_NUMERIC \
+  (100000 * MLX_VERSION_MAJOR + 1000 * MLX_VERSION_MINOR + MLX_VERSION_PATCH)
+
+namespace mlx::core {
+
+/* A string representation of the MLX version in the format
+ * "major.minor.patch".
+ *
+ * For dev builds, the version will include the suffix ".devYYYYMMDD+hash"
+ */
+std::string version();
+
+} // namespace mlx::core

python/mlx/distributed_run.py

@@ -14,8 +14,11 @@ import time
 from collections import Counter
 from dataclasses import dataclass
 from pathlib import Path
+from queue import Empty as QueueEmpty
+from queue import Queue
 from select import select
 from subprocess import PIPE, Popen, run
+from typing import Optional
 
 
 @dataclass
@@ -25,6 +28,100 @@ class Host:
     ips: list[str]
 
 
+@dataclass
+class ThunderboltPort:
+    iface: str
+    uuid: str
+    connected_to: Optional[str]
+
+
+@dataclass
+class ThunderboltHost:
+    name: str
+    ports: list[ThunderboltPort]
+
+
+def parse_hardware_ports(ports_string):
+    ports = {}
+    port_name = None
+    for l in ports_string.decode("utf-8").split("\n"):
+        if l.startswith("Hardware Port:"):
+            port_name = l.strip()[15:]
+        elif l.startswith("Device:"):
+            ports[port_name] = l.strip()[8:]
+            port_name = None
+    return ports
+
+
+def extract_rings(hosts, index):
+    def usable_port(i, j, used_ports):
+        return (i, j) not in used_ports and hosts[i].ports[j].connected_to is not None
+
+    def dfs(start_node, node, path, visited, used_ports):
+        path.append(node)
+        visited.add(node)
+        for j, p in enumerate(hosts[node].ports):
+            if not usable_port(node, j, used_ports):
+                continue
+            next_node, _ = index[p.connected_to]
+            if next_node == start_node:
+                yield path[:]
+            if next_node not in visited:
+                yield from dfs(start_node, next_node, path, visited, used_ports)
+        path.pop()
+        visited.remove(node)
+
+    # Concretize maps the found cycle to real thunderbolt ports. It also adds
+    # those ports to the used set so next cycles can't use them again.
+    def concretize(cycle, used_ports):
+        concrete_path = []
+        for n1, n2 in zip(cycle, cycle[1:] + cycle[:1]):
+            for j, p in enumerate(hosts[n1].ports):
+                if not usable_port(n1, j, used_ports):
+                    continue
+                n2_hat, nj = index[p.connected_to]
+                if n2 == n2_hat:
+                    concrete_path.append(((n1, j), (n2, nj)))
+                    used_ports.add((n1, j))
+                    used_ports.add((n2, nj))
+                    break
+            if concrete_path[-1][0][0] != n1:
+                raise RuntimeError("Couldn't concretize the cycle")
+        return concrete_path
+
+    # Normalize tries to ensure that the cycles have the same direction so we can
+    # use them together. We achieve this by selecting the direction such that
+    # the smallest rank hosts connect to larger rank hosts.
+    def normalize(path):
+        small_to_large = sum(1 for p in path if p[0][0] < p[1][0])
+        if small_to_large > len(path) - small_to_large:
+            return path
+        else:
+            return [(p[1], p[0]) for p in path]
+
+    rings = []
+    used_ports = set()
+    for start_node in range(len(hosts)):
+        while True:
+            ring = []
+            for r in dfs(start_node, start_node, [], set(), used_ports):
+                if len(r) > len(ring):
+                    ring = r
+                # Break early since we won't find a bigger ring no matter what
+                if len(ring) == len(hosts):
+                    break
+            if not ring:
+                break
+            try:
+                rings.append(normalize(concretize(ring, used_ports)))
+            except RuntimeError:
+                if len(rings) > 0:
+                    return rings
+                raise
+
+    return rings
+
+
 def positive_number(x):
     x = int(x)
     if x <= 0:
@@ -43,6 +140,11 @@ def log_warning(*args, **kwargs):
     print("\033[33m[WARN]", *args, "\033[0m", **kwargs)
 
 
+def log_error(*args, **kwargs):
+    kwargs["file"] = sys.stderr
+    print("\033[31m[ERROR]", *args, "\033[0m", **kwargs)
+
+
 def parse_hostfile(parser, hostfile):
     """Parse the json hostfile that contains both the hostnames to ssh into and
     the ips to communicate over when using the ring backend.
@@ -77,6 +179,8 @@ def parse_hostfile(parser, hostfile):
 def parse_hostlist(parser, hostlist, repeats):
     hosts = []
     for i, h in enumerate(hostlist.split(",")):
+        if h == "":
+            raise ValueError("Hostname cannot be empty")
         try:
             ipaddress.ip_address(h)
             ips = [h]
@@ -88,46 +192,54 @@ def parse_hostlist(parser, hostlist, repeats):
 
 
 def make_monitor_script(rank, hostfile, cwd, env, command, verbose):
+    # Imports that are used throughout
     script = ""
+    script += "import os\n"
+    script += "import sys\n"
+    script += "import tempfile\n"
+    script += "from pathlib import Path\n"
 
     # Write the PID to a file so we can kill the process if needed
-    script += "pidfile=$(mktemp)\n"
-    script += "echo $$ >$pidfile\n"
-    script += "echo $pidfile\n"
+    script += "_, pidfile = tempfile.mkstemp() \n"
+    script += "open(pidfile, 'w').write(str(os.getpid()))\n"
+    script += "print(pidfile, flush=True)\n"
 
     # Change the working directory if one was requested. Otherwise attempt to
-    # change to change to the current one but don't fail if it wasn't possible.
+    # change to the current one but don't fail if it wasn't possible.
     d = cwd or os.getcwd()
-    script += f"if [ -d {shlex.quote(d)} ]; then\n"
-    script += f"    cd {shlex.quote(d)}\n"
+    script += f"if Path({repr(d)}).exists():\n"
+    script += f"    os.chdir({repr(d)})\n"
     if cwd is not None:
-        script += "else\n"
-        script += f"    echo Failed to change directory to {shlex.quote(d)} 1>&2\n"
-        script += f"    exit 1\n"
-    script += "fi\n"
+        script += "else:\n"
+        script += (
+            f"    print('Failed to change directory to', {repr(d)}, file=sys.stderr)\n"
+        )
+        script += f"    sys.exit(1)\n"
 
     # Add the environment variables that were given to us
+    script += "env = dict(os.environ)\n"
     for e in env:
         key, *value = e.split("=", maxsplit=1)
         value = shlex.quote(value[0]) if len(value) > 0 else ""
         if not all(c.isalnum() or c == "_" for c in key):
             log_warning(f"'{e}' is an invalid environment variable so it is ignored")
             continue
-        script += f"export {key}={value}\n"
+        script += f"env[{repr(key)}] = {repr(value)}\n"
 
     # Add the environment variables to enable the ring distributed backend
     if hostfile != "":
-        script += "tmpfile=$(mktemp)\n"
-        script += f"echo {shlex.quote(hostfile)} >$tmpfile\n"
+        script += "_, hostfile = tempfile.mkstemp()\n"
+        script += "with open(hostfile, 'w') as f:\n"
+        script += f"    f.write({repr(hostfile)})\n"
         if verbose:
-            script += "export MLX_RING_VERBOSE=1\n"
-        script += "export MLX_HOSTFILE=$tmpfile\n"
-        script += f"export MLX_RANK={rank}\n"
+            script += "env['MLX_RING_VERBOSE'] = '1'\n"
+        script += "env['MLX_HOSTFILE'] = hostfile\n"
+        script += f"env['MLX_RANK'] = '{rank}'\n"
         script += "\n"
 
     # Replace the process with the script
-    script += shlex.join(["exec", *command])
-    script += "\n"
+    script += f"command = [{','.join(map(repr, command))}]\n"
+    script += "os.execve(command[0], command, env)\n"
 
     return script
 
@@ -136,28 +248,37 @@ def launch_ring(parser, hosts, args, command):
     stop = False
     exit_codes = [None] * len(hosts)
 
-    def node_thread(rank, host, hostfile):
+    def node_thread(rank, host, hostfile, input_queue):
         is_local = host == "127.0.0.1"
         script = make_monitor_script(
             rank, hostfile, args.cwd, args.env, command, args.verbose
         )
         script_b64 = base64.b64encode(script.encode()).decode()
-        cmd = f'echo "{script_b64}" | base64 -d | /bin/bash'
+        cmd = f'{sys.executable} -c "import base64; exec(base64.b64decode(\\"{script_b64}\\"));"'
         if not is_local:
             cmd = f"ssh {host} '{cmd}'"
         p = Popen(
             cmd,
             shell=True,
+            stdin=PIPE,
             stdout=PIPE,
             stderr=PIPE,
         )
         os.set_blocking(p.stdout.fileno(), False)
         os.set_blocking(p.stderr.fileno(), False)
+        os.set_blocking(p.stdin.fileno(), False)
 
         # Repeat the stdout and stderr to the local machine
+        to_read = [p.stdout.fileno(), p.stderr.fileno()]
+        to_write = [p.stdin.fileno()]
         pidfile = ""
+        stdin_buffer = b""
         while p.poll() is None:
-            rlist, _, _ = select([p.stdout.fileno(), p.stderr.fileno()], [], [], 1.0)
+            try:
+                stdin_buffer += input_queue.get_nowait()
+            except QueueEmpty:
+                pass
+            rlist, wlist, _ = select(to_read, to_write, [], 1.0)
             for fd in rlist:
                 is_stdout = fd == p.stdout.fileno()
                 outfile = sys.stdout if is_stdout else sys.stderr
@@ -169,6 +290,11 @@ def launch_ring(parser, hosts, args, command):
                     msg = msg[0] if msg else ""
 
                 outfile.write(msg)
+                outfile.flush()
+            for fd in wlist:
+                if len(stdin_buffer) > 0:
+                    n = os.write(fd, stdin_buffer)
+                    stdin_buffer = stdin_buffer[n:]
             if stop:
                 p.terminate()
                 break
@@ -200,7 +326,7 @@ def launch_ring(parser, hosts, args, command):
             "The ring backend requires IPs to be provided instead of hostnames"
         )
 
-    port = 5000
+    port = args.starting_port
     ring_hosts = []
     for h in hosts:
         node = []
@@ -213,16 +339,25 @@ def launch_ring(parser, hosts, args, command):
 
     log(args.verbose, "Running", shlex.join(command))
 
+    input_queues = []
     threads = []
     for i, h in enumerate(hosts):
         if i + 1 == len(hosts):
             time.sleep(1.0)
-        t = threading.Thread(target=node_thread, args=(i, h.ssh_hostname, hostfile))
+        input_queues.append(Queue())
+        t = threading.Thread(
+            target=node_thread, args=(i, h.ssh_hostname, hostfile, input_queues[-1])
+        )
         t.start()
         threads.append(t)
 
+    os.set_blocking(sys.stdin.fileno(), False)
     while not stop:
-        time.sleep(1.0)
+        rlist, _, _ = select([sys.stdin.fileno()], [], [], 1.0)
+        for fd in rlist:
+            stdin_buffer = os.read(fd, 8192)
+            for q in input_queues:
+                q.put(stdin_buffer)
         if any(t.is_alive() for t in threads):
             for i, t in enumerate(threads):
                 if not t.is_alive():
@@ -271,8 +406,312 @@ def launch_mpi(parser, hosts, args, command):
             pass
 
 
+def check_ssh_connections(hosts):
+    results = [False] * len(hosts)
+
+    def _check(hostname, i):
+        result = run(
+            [
+                "ssh",
+                "-o",
+                "BatchMode=yes",
+                "-o",
+                "ConnectTimeout=5",
+                hostname,
+                "echo",
+                "success",
+            ],
+            stdout=PIPE,
+            stderr=PIPE,
+        )
+        results[i] = result.returncode == 0
+
+    threads = [
+        threading.Thread(target=_check, args=(h.ssh_hostname, i))
+        for i, h in enumerate(hosts)
+    ]
+    for t in threads:
+        t.start()
+    for t in threads:
+        t.join()
+
+    if not all(results):
+        log_error("Could not ssh to the following hosts:")
+        for i, h in enumerate(hosts):
+            if not results[i]:
+                log_error("  - ", h.ssh_hostname)
+        log_error()
+        log_error("Maybe they are not set-up for password-less ssh?")
+        sys.exit(1)
+
+
+def prepare_tb_ring(args, hosts):
+    log(
+        args.verbose,
+        f"Preparing a thunderbolt ring for {', '.join(h.ssh_hostname for h in hosts)}",
+    )
+
+    # Check that we can ssh
+    check_ssh_connections(hosts)
+    if args.auto_setup and args.verbose:
+        log_warning(
+            "--auto-setup is requested which requires password-less sudo",
+            "on the remote hosts",
+        )
+
+    # Extract the current connectivity from the remote hosts
+    thunderbolt_connections = []
+    for h in hosts:
+        log(args.verbose, "Getting connectivity from", h.ssh_hostname)
+        thunderbolt_connections.append(
+            json.loads(
+                run(
+                    [
+                        "ssh",
+                        h.ssh_hostname,
+                        "system_profiler",
+                        "SPThunderboltDataType",
+                        "-json",
+                    ],
+                    capture_output=True,
+                ).stdout
+            )
+        )
+    interface_maps = []
+    for h in hosts:
+        log(args.verbose, "Getting interface names from", h.ssh_hostname)
+        interface_maps.append(
+            parse_hardware_ports(
+                run(
+                    [
+                        "ssh",
+                        h.ssh_hostname,
+                        "networksetup",
+                        "-listallhardwareports",
+                    ],
+                    capture_output=True,
+                ).stdout
+            )
+        )
+
+    # Parse the connectivity into some simple dataclasses
+    tb_hosts = []
+    for c, iface_map in zip(thunderbolt_connections, interface_maps):
+        name = ""
+        ports = []
+        for t in c["SPThunderboltDataType"]:
+            name = t["device_name_key"]
+            uuid = t["domain_uuid_key"]
+            tag = t["receptacle_1_tag"]["receptacle_id_key"]
+            if items := t.get("_items", []):
+                connected_to = items[0]["domain_uuid_key"]
+            else:
+                connected_to = None
+            iface = iface_map[f"Thunderbolt {tag}"]
+            ports.append(ThunderboltPort(iface, uuid, connected_to))
+        tb_hosts.append(ThunderboltHost(name, sorted(ports, key=lambda x: x.iface)))
+
+    # Create a reverse index to be able to map uuids to (host, port) quickly
+    uuid_reverse_index = {}
+    for i, h in enumerate(tb_hosts):
+        for j, p in enumerate(h.ports):
+            uuid_reverse_index[p.uuid] = (i, j)
+
+    # Find the rings by simply walking and marking visited (host, port) tuples
+    # and keeping the largest rings greedily.
+    log(args.verbose, "Extracting rings from the parsed connectivity")
+    rings = extract_rings(tb_hosts, uuid_reverse_index)
+
+    # Just output a DOT graphical representation of the found rings
+    if args.dot:
+        names = []
+        for i in range(len(tb_hosts)):
+            n = ""
+            j = i
+            while True:
+                n += chr(97 + j % 26)
+                j //= 26
+                if j == 0:
+                    break
+            names.append(n)
+
+        print("graph G {")
+        print("  node [shape=rectangle];")
+        for i, h in enumerate(hosts):
+            print(f'  {names[i]} [label="{h.ssh_hostname}"];')
+        for r in rings:
+            for (i, _), (j, _) in r:
+                print(f"  {names[i]} -- {names[j]};")
+        print("}")
+        return
+
+    # Assign IPs to each interface such that the interfaces can communicate
+    ips = {}
+    pairs = {}
+    expecting = set()
+    ip0 = 0
+    ip1 = 0
+    netmask = "255.255.255.252"
+    for r in rings:
+        for a, b in r:
+            ips[a] = f"192.168.{ip0}.{ip1 + 1}"
+            ips[b] = f"192.168.{ip0}.{ip1 + 2}"
+            pairs[a] = b
+            pairs[b] = a
+            expecting.add(b)
+            ip1 += 4
+            if ip1 > 255:
+                ip0 += 1
+                ip1 = 0
+            if ip0 > 255:
+                raise ValueError("Ran out of available local IPs for the ring")
+
+    # Create the hostfile
+    hostfile = []
+    for i, h in enumerate(hosts):
+        host = {
+            "ssh": h.ssh_hostname,
+            "ips": [
+                ips[i, j]
+                for j, p in enumerate(tb_hosts[i].ports)
+                if (i, j) in expecting
+            ],
+        }
+        hostfile.append(host)
+
+    if not args.hostfile_only:
+        for i, h in enumerate(hosts):
+            command = ""
+            command += "sudo ifconfig bridge0 down\n"
+            for j, p in enumerate(tb_hosts[i].ports):
+                if (i, j) not in ips:
+                    continue
+                iface = p.iface
+                ip = ips[i, j]
+                peer = ips[pairs[i, j]]
+                command += f"sudo ifconfig {iface} inet {ip} netmask {netmask}\n"
+                command += f"sudo route change {peer} -interface {iface}\n"
+            if args.auto_setup:
+                print(f"Running auto setup for {h.ssh_hostname}")
+                command = command.strip().replace("\n", " && ")
+                command = ["ssh", h.ssh_hostname, command]
+                log(args.verbose, shlex.join(command))
+                run(command)
+            else:
+                msg = f"Setup for {h.ssh_hostname}"
+                print(msg)
+                print("=" * len(msg))
+                print(command)
+                input("Enter to continue")
+            print()
+
+    if args.output_hostfile:
+        with open(args.output_hostfile, "w") as f:
+            json.dump(hostfile, f, indent=4)
+    else:
+        print("Hostfile")
+        print("========")
+        print(json.dumps(hostfile, indent=4))
+
+
+def prepare_hostfile(args, hosts):
+    log(
+        args.verbose,
+        f"Preparing an ethernet hostfile for {', '.join(h.ssh_hostname for h in hosts)}",
+    )
+
+    # Check that we can ssh
+    check_ssh_connections(hosts)
+
+    # Get the ips for each host
+    for h in hosts:
+        log(args.verbose, "Getting the ip from", h.ssh_hostname)
+        h.ips.append(
+            run(
+                ["ssh", h.ssh_hostname, "ipconfig", "getifaddr", "en0"],
+                capture_output=True,
+                text=True,
+            ).stdout.strip()
+        )
+
+    hostfile = []
+    for h in hosts:
+        hostfile.append(dict(ssh=h.ssh_hostname, ips=h.ips))
+
+    if args.output_hostfile:
+        with open(args.output_hostfile, "w") as f:
+            json.dump(hostfile, f, indent=4)
+    else:
+        print("Hostfile")
+        print("========")
+        print(json.dumps(hostfile, indent=4))
+
+
+def distributed_config():
+    parser = argparse.ArgumentParser(
+        description="Configure remote machines for use with MLX distributed"
+    )
+    parser.add_argument(
+        "--verbose", action="store_true", help="Print debug messages in stdout"
+    )
+    parser.add_argument(
+        "--backend",
+        choices=["ring", "mpi"],
+        default="ring",
+        help="Which distributed backend to configure",
+    )
+    parser.add_argument(
+        "--over",
+        choices=["thunderbolt", "ethernet"],
+        default="thunderbolt",
+        help="What type of connectivity to configure",
+    )
+    parser.add_argument(
+        "--hosts", default="127.0.0.1", help="A comma separated list of hosts"
+    )
+    parser.add_argument("--hostfile", help="The file containing the hosts")
+    parser.add_argument(
+        "--dot", action="store_true", help="Output the topology in DOT format and exit"
+    )
+    parser.add_argument(
+        "--hostfile-only", action="store_true", help="If set only compute the hostfile"
+    )
+    parser.add_argument(
+        "--output-hostfile", help="If provided, save the hostfile to this path"
+    )
+    parser.add_argument(
+        "--auto-setup",
+        action="store_true",
+        help="If set we will attempt to automatically configure the machines via ssh",
+    )
+    args = parser.parse_args()
+
+    if args.backend == "mpi" and args.over == "thunderbolt":
+        raise ValueError(
+            (
+                "The configuration of MPI over thunderbolt is "
+                "not supported yet by mlx.distributed_config"
+            )
+        )
+
+    if args.hostfile is not None:
+        hosts = parse_hostfile(parser, args.hostfile)
+    else:
+        hosts = parse_hostlist(parser, args.hosts, 1)
+
+    if args.over == "thunderbolt":
+        prepare_tb_ring(args, hosts)
+    else:
+        prepare_hostfile(args, hosts)
+
+
 def main():
     parser = argparse.ArgumentParser(description="Launch an MLX distributed program")
+    parser.add_argument(
+        "--print-python",
+        action="store_true",
+        help="Print the path to the current python executable and exit",
+    )
     parser.add_argument(
         "--verbose", action="store_true", help="Print debug messages in stdout"
     )
@@ -311,11 +750,27 @@ def main():
         type=int,
         help="How many connections per ip to use for the ring backend",
     )
+    parser.add_argument(
+        "--starting-port",
+        "-p",
+        type=int,
+        default=5000,
+        help="For the ring backend listen on this port increasing by 1 per rank and IP",
+    )
     parser.add_argument(
         "--cwd", help="Set the working directory on each node to the provided one"
     )
     args, rest = parser.parse_known_args()
 
+    if args.print_python:
+        print(sys.executable)
+        return
+
+    if len(rest) == 0:
+        parser.error("No script is provided")
+    if rest[0] == "--":
+        rest.pop(0)
+
     # Try to extract a list of hosts and corresponding ips
     if args.hostfile is not None:
         hosts = parse_hostfile(parser, args.hostfile)

python/mlx/nn/__init__.py

@@ -2,4 +2,4 @@
 
 from mlx.nn import init, losses
 from mlx.nn.layers import *
-from mlx.nn.utils import value_and_grad
+from mlx.nn.utils import average_gradients, value_and_grad

python/mlx/optimizers/optimizers.py

@@ -5,7 +5,7 @@ from typing import Callable, List, Optional, Tuple, Union
 
 import mlx.core as mx
 from mlx.nn import Module
-from mlx.utils import tree_map, tree_reduce
+from mlx.utils import tree_flatten, tree_map, tree_merge, tree_reduce, tree_unflatten
 
 
 class Optimizer:
@@ -154,6 +154,79 @@ class Optimizer:
         self.state[name] = param
 
 
+class MultiOptimizer(Optimizer):
+    """Wraps a list of optimizers with corresponding weight predicates/filters
+    to make it easy to use different optimizers for different weights.
+
+    The predicates take the full "path" of the weight and the weight itself and
+    return True if it should be considered for this optimizer. The last
+    optimizer in the list is a fallback optimizer and no predicate should be
+    given for it.
+
+    Args:
+        optimizers (list[Optimizer]): A list of optimizers to delegate to
+        filters (list[Callable[[str, array], bool]): A list of predicates that
+            should be one less than the provided optimizers.
+    """
+
+    def __init__(self, optimizers, filters: list = []):
+        super().__init__()
+        self._state = {}
+
+        if len(filters) != len(optimizers) - 1:
+            raise ValueError(
+                f"Given {len(filters)} filters but {len(optimizers)-1} needed."
+            )
+
+        self.optimizers = optimizers
+        self.filters = filters + [lambda *args, **kwargs: True]
+
+    def _split_dictionary(self, gradients: dict):
+        if len(self.optimizers) == 1:
+            return [gradients]
+
+        parts = [[] for _ in range(len(self.optimizers))]
+        flat_gradients = tree_flatten(gradients)
+        for k, g in flat_gradients:
+            for i, fn in enumerate(self.filters):
+                if fn(k, g):
+                    parts[i].append((k, g))
+                    break
+
+        return [tree_unflatten(p) for p in parts]
+
+    def init(self, parameters: dict):
+        for o, p in zip(self.optimizers, self._split_dictionary(parameters)):
+            o.init(p)
+
+    def apply_gradients(self, gradients: dict, parameters: dict):
+        tree = {}
+        for o, g in zip(self.optimizers, self._split_dictionary(gradients)):
+            tree = tree_merge(tree, o.apply_gradients(g, parameters))
+        return tree
+
+    @property
+    def state(self):
+        return {"states": [o.state for o in self.optimizers]}
+
+    @state.setter
+    def state(self, state: dict):
+        if "states" not in state or len(state["states"]) != len(self.optimizers):
+            raise ValueError("Invalid state provided")
+
+        for o, s in zip(self.optimizers, state["states"]):
+            o.state = s
+
+    @property
+    def learning_rate(self):
+        return self.optimizers[0].learning_rate
+
+    @learning_rate.setter
+    def learning_rate(self, learning_rate: Union[float, mx.array]):
+        for o in self.optimizers:
+            o.learning_rate = learning_rate
+
+
 class SGD(Optimizer):
     r"""The stochastic gradient descent optimizer.
 

python/mlx/utils.py

@@ -1,5 +1,6 @@
 # Copyright © 2023 Apple Inc.
 from collections import defaultdict
+from itertools import zip_longest
 from typing import Any, Callable, List, Optional, Tuple
 
 
@@ -244,3 +245,46 @@ def tree_reduce(fn, tree, initializer=None, is_leaf=None):
         return tree if accumulator is None else fn(accumulator, tree)
 
     return accumulator
+
+
+def tree_merge(tree_a, tree_b, merge_fn=None):
+    """Merge two Python trees in one containing the values of both. It can be
+    thought of as a deep dict.update method.
+
+    Args:
+        tree_a (Any): The first Python tree.
+        tree_b (Any): The second Python tree.
+        merge_fn (callable, optional): A function to merge leaves.
+
+    Returns:
+        The Python tree containing the values of both ``tree_a`` and
+        ``tree_b``.
+    """
+    if isinstance(tree_a, (dict, list, tuple)) and len(tree_a) == 0:
+        tree_a = None
+    if isinstance(tree_b, (dict, list, tuple)) and len(tree_b) == 0:
+        tree_b = None
+    if tree_a is None and tree_b is not None:
+        return tree_b
+    if tree_a is not None and tree_b is None:
+        return tree_a
+
+    if isinstance(tree_a, (list, tuple)) and isinstance(tree_b, (list, tuple)):
+        TreeType = type(tree_a)
+        return TreeType(
+            tree_merge(a, b, merge_fn) for a, b in zip_longest(tree_a, tree_b)
+        )
+    elif isinstance(tree_a, dict) and isinstance(tree_b, dict):
+        return {
+            k: tree_merge(tree_a.get(k, None), tree_b.get(k, None), merge_fn)
+            for k in set(tree_a.keys()) | set(tree_b.keys())
+        }
+    else:
+        if merge_fn is None:
+            raise ValueError(
+                (
+                    "Trees contain elements at the same locations but no merge "
+                    "function was provided"
+                )
+            )
+        return merge_fn(tree_a, tree_b)

python/src/array.cpp

@@ -878,6 +878,38 @@ void init_array(nb::module_& m) {
           },
           "other"_a,
           nb::rv_policy::none)
+      .def(
+          "__xor__",
+          [](const mx::array& a, const ScalarOrArray v) {
+            if (!is_comparable_with_array(v)) {
+              throw_invalid_operation("bitwise xor", v);
+            }
+            auto b = to_array(v, a.dtype());
+            if (mx::issubdtype(a.dtype(), mx::inexact) ||
+                mx::issubdtype(b.dtype(), mx::inexact)) {
+              throw std::invalid_argument(
+                  "Floating point types not allowed with bitwise xor.");
+            }
+            return mx::bitwise_xor(a, b);
+          },
+          "other"_a)
+      .def(
+          "__ixor__",
+          [](mx::array& a, const ScalarOrArray v) -> mx::array& {
+            if (!is_comparable_with_array(v)) {
+              throw_invalid_operation("inplace bitwise xor", v);
+            }
+            auto b = to_array(v, a.dtype());
+            if (mx::issubdtype(a.dtype(), mx::inexact) ||
+                mx::issubdtype(b.dtype(), mx::inexact)) {
+              throw std::invalid_argument(
+                  "Floating point types not allowed bitwise xor.");
+            }
+            a.overwrite_descriptor(mx::bitwise_xor(a, b));
+            return a;
+          },
+          "other"_a,
+          nb::rv_policy::none)
       .def("__int__", [](mx::array& a) { return nb::int_(to_scalar(a)); })
       .def("__float__", [](mx::array& a) { return nb::float_(to_scalar(a)); })
       .def(

python/src/buffer.h

@@ -44,6 +44,8 @@ std::string buffer_format(const mx::array& a) {
       return "f";
     case mx::bfloat16:
       return "B";
+    case mx::float64:
+      return "d";
     case mx::complex64:
       return "Zf\0";
     default: {

python/src/convert.cpp

@@ -152,6 +152,8 @@ nb::ndarray<NDParams...> mlx_to_nd_array(const mx::array& a) {
       throw nb::type_error("bfloat16 arrays cannot be converted to NumPy.");
     case mx::float32:
       return mlx_to_nd_array_impl<float, NDParams...>(a);
+    case mx::float64:
+      return mlx_to_nd_array_impl<double, NDParams...>(a);
     case mx::complex64:
       return mlx_to_nd_array_impl<std::complex<float>, NDParams...>(a);
     default:

python/src/distributed.cpp

@@ -10,6 +10,8 @@
 #include "mlx/distributed/distributed.h"
 #include "mlx/distributed/ops.h"
 
+#include "python/src/utils.h"
+
 namespace mx = mlx::core;
 namespace nb = nanobind;
 using namespace nb::literals;
@@ -66,19 +68,21 @@ void init_distributed(nb::module_& parent_module) {
 
         Example:
 
-          import mlx.core as mx
+          .. code:: python
 
-          group = mx.distributed.init(backend="ring")
+            import mlx.core as mx
 
+            group = mx.distributed.init(backend="ring")
 
         Args:
           strict (bool, optional): If set to False it returns a singleton group
             in case ``mx.distributed.is_available()`` returns False otherwise
             it throws a runtime error. Default: ``False``
-          backend (str, optional): Select a specific distributed backend to
-            initialize. If set to ``any`` then try all available backends and
-            return the first one that succeeds. Subsequent calls will return
-            the first backend that was initialized. Default: ``any``
+          backend (str, optional): Which distributed backend to initialize.
+            Possible values ``mpi``, ``ring``, ``any``. If set to ``any`` all
+            available backends are tried and the first one that succeeds
+            becomes the global group which will be returned in subsequent
+            calls. Default: ``any``
 
         Returns:
           Group: The group representing all the launched processes.
@@ -86,7 +90,11 @@ void init_distributed(nb::module_& parent_module) {
 
   m.def(
       "all_sum",
-      &mx::distributed::all_sum,
+      [](const ScalarOrArray& x,
+         std::optional<mx::distributed::Group> group,
+         mx::StreamOrDevice s) {
+        return mx::distributed::all_sum(to_array(x), group, s);
+      },
       "x"_a,
       nb::kw_only(),
       "group"_a = nb::none(),
@@ -112,7 +120,11 @@ void init_distributed(nb::module_& parent_module) {
 
   m.def(
       "all_gather",
-      &mx::distributed::all_gather,
+      [](const ScalarOrArray& x,
+         std::optional<mx::distributed::Group> group,
+         mx::StreamOrDevice s) {
+        return mx::distributed::all_gather(to_array(x), group, s);
+      },
       "x"_a,
       nb::kw_only(),
       "group"_a = nb::none(),
@@ -139,7 +151,12 @@ void init_distributed(nb::module_& parent_module) {
 
   m.def(
       "send",
-      &mx::distributed::send,
+      [](const ScalarOrArray& x,
+         int dst,
+         std::optional<mx::distributed::Group> group,
+         mx::StreamOrDevice s) {
+        return mx::distributed::send(to_array(x), dst, group, s);
+      },
       "x"_a,
       "dst"_a,
       nb::kw_only(),
@@ -195,7 +212,12 @@ void init_distributed(nb::module_& parent_module) {
 
   m.def(
       "recv_like",
-      &mx::distributed::recv_like,
+      [](const ScalarOrArray& x,
+         int src,
+         std::optional<mx::distributed::Group> group,
+         mx::StreamOrDevice s) {
+        return mx::distributed::recv_like(to_array(x), src, group, s);
+      },
       "x"_a,
       "src"_a,
       nb::kw_only(),

python/src/fast.cpp

@@ -25,12 +25,12 @@ void init_fast(nb::module_& parent_module) {
       "rms_norm",
       &mx::fast::rms_norm,
       "x"_a,
-      "weight"_a,
+      "weight"_a.none(),
       "eps"_a,
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def rms_norm(x: array, weight: array, eps: float, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def rms_norm(x: array, weight: Optional[array], eps: float, *, stream: Union[None, Stream, Device] = None) -> array"),
       R"pbdoc(
         Root Mean Square normalization (RMS norm).
 
@@ -38,9 +38,9 @@ void init_fast(nb::module_& parent_module) {
 
         Args:
             x (array): Input array.
-            weight (array): A multiplicative weight to scale the result by.
+            weight (array, optional): A multiplicative weight to scale the result by.
               The ``weight`` should be one-dimensional with the same size
-              as the last axis of ``x``.
+              as the last axis of ``x``. If set to ``None`` then no scaling happens.
             eps (float): A small additive constant for numerical stability.
 
         Returns:

python/src/linalg.cpp

@@ -92,6 +92,7 @@ void init_linalg(nb::module_& parent_module) {
           =====  ============================  ==========================
           None   Frobenius norm                2-norm
           'fro'  Frobenius norm                --
+          'nuc'  nuclear norm                  --
           inf    max(sum(abs(x), axis=1))      max(abs(x))
           -inf   min(sum(abs(x), axis=1))      min(abs(x))
           0      --                            sum(x != 0)
@@ -102,9 +103,6 @@ void init_linalg(nb::module_& parent_module) {
           other  --                            sum(abs(x)**ord)**(1./ord)
           =====  ============================  ==========================
 
-          .. warning::
-            Nuclear norm and norms based on singular values are not yet implemented.
-
           The Frobenius norm is given by [1]_:
 
               :math:`||A||_F = [\sum_{i,j} abs(a_{i,j})^2]^{1/2}`
@@ -206,15 +204,22 @@ void init_linalg(nb::module_& parent_module) {
       )pbdoc");
   m.def(
       "svd",
-      [](const mx::array& a, mx::StreamOrDevice s /* = {} */) {
-        const auto result = mx::linalg::svd(a, s);
-        return nb::make_tuple(result.at(0), result.at(1), result.at(2));
+      [](const mx::array& a,
+         bool compute_uv /* = true */,
+         mx::StreamOrDevice s /* = {} */) -> nb::object {
+        const auto result = mx::linalg::svd(a, compute_uv, s);
+        if (result.size() == 1) {
+          return nb::cast(result.at(0));
+        } else {
+          return nb::make_tuple(result.at(0), result.at(1), result.at(2));
+        }
       },
       "a"_a,
+      "compute_uv"_a = true,
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def svd(a: array, *, stream: Union[None, Stream, Device] = None) -> Tuple[array, array, array]"),
+          "def svd(a: array, compute_uv: bool = True, *, stream: Union[None, Stream, Device] = None) -> Tuple[array, array, array]"),
       R"pbdoc(
         The Singular Value Decomposition (SVD) of the input matrix.
 
@@ -224,12 +229,15 @@ void init_linalg(nb::module_& parent_module) {
 
         Args:
             a (array): Input array.
+            compute_uv (bool, optional): If ``True``, return the ``U``, ``S``, and ``Vt`` components.
+              If ``False``, return only the ``S`` array. Default: ``True``.
             stream (Stream, optional): Stream or device. Defaults to ``None``
               in which case the default stream of the default device is used.
 
         Returns:
-            tuple(array, array, array): The ``U``, ``S``, and ``Vt`` matrices, such that
-            ``A = U @ diag(S) @ Vt``
+            Union[tuple(array, ...), array]:
+              If compute_uv is ``True`` returns the ``U``, ``S``, and ``Vt`` matrices, such that 
+              ``A = U @ diag(S) @ Vt``. If compute_uv is ``False`` returns singular values array ``S``.
       )pbdoc");
   m.def(
       "inv",

python/tests/mpi_test_distributed.py

@@ -177,6 +177,7 @@ class TestDistributed(mlx_tests.MLXTestCase):
     def test_donation(self):
         x = mx.random.normal((1024,))
         mx.eval(x)
+        mx.synchronize(mx.default_stream(mx.default_device()))
 
         mx.metal.reset_peak_memory()
         scale = mx.array(2.0)

python/tests/ring_test_distributed.py

@@ -56,6 +56,45 @@ class TestRingDistributed(mlx_tests.MLXTestCase):
                 maxrelerror = ((y - z).abs() / z.abs()).max()
                 self.assertLessEqual(maxrelerror, rtol)
 
+    def test_send_recv(self):
+        world = mx.distributed.init()
+        dtypes = [
+            mx.int8,
+            mx.uint8,
+            mx.int16,
+            mx.uint16,
+            mx.int32,
+            mx.uint32,
+            mx.float32,
+            mx.float16,
+            mx.bfloat16,
+            mx.complex64,
+        ]
+        sizes = [
+            (7,),
+            (10,),
+            (1024,),
+            (1024, 1024),
+        ]
+        key = mx.random.key(0)
+        right = (world.rank() + 1) % world.size()
+        left = (world.rank() + world.size() - 1) % world.size()
+        for dt in dtypes:
+            for sh in sizes:
+                x = (
+                    mx.random.uniform(shape=(world.size(),) + sh, key=key) * 10
+                ).astype(dt)
+                if world.rank() % 2 == 0:
+                    y = mx.distributed.send(x[world.rank()], right)
+                    z = mx.distributed.recv_like(y, left)
+                    mx.eval(y, z)
+                else:
+                    z = mx.distributed.recv_like(x[world.rank()], left)
+                    y = mx.distributed.send(x[world.rank()], right)
+                    mx.eval(z, y)
+                self.assertTrue(mx.all(y == x[world.rank()]))
+                self.assertTrue(mx.all(z == x[left]))
+
 
 if __name__ == "__main__":
     unittest.main()

python/tests/test_array.py

@@ -1725,6 +1725,7 @@ class TestArray(mlx_tests.MLXTestCase):
         self.assertEqual((mx.array(True) | False).item(), True)
         self.assertEqual((mx.array(False) | False).item(), False)
         self.assertEqual((~mx.array(False)).item(), True)
+        self.assertEqual((mx.array(False) ^ True).item(), True)
 
     def test_inplace(self):
         iops = [
@@ -1734,6 +1735,7 @@ class TestArray(mlx_tests.MLXTestCase):
             "__ifloordiv__",
             "__imod__",
             "__ipow__",
+            "__ixor__",
         ]
 
         for op in iops:
@@ -1773,6 +1775,10 @@ class TestArray(mlx_tests.MLXTestCase):
         b @= a
         self.assertTrue(mx.array_equal(a, b))
 
+        a = mx.array(False)
+        a ^= True
+        self.assertEqual(a.item(), True)
+
     def test_inplace_preserves_ids(self):
         a = mx.array([1.0])
         orig_id = id(a)

python/tests/test_compile.py

@@ -815,6 +815,31 @@ class TestCompile(mlx_tests.MLXTestCase):
         out = fun(*inputs)
         self.assertTrue(mx.allclose(out, mx.full((2, 2), 20)))
 
+        @mx.compile
+        def fun(arrs):
+            for _ in range(6):
+                arrs = [x + y for x, y in zip(arrs[::2], arrs[1::2])]
+            return arrs[0]
+
+        arrs = [mx.array([1.0, 2.0]) for _ in range(64)]
+        out = fun(arrs)
+        self.assertTrue(mx.allclose(out, mx.array([64.0, 128.0])))
+
+    def test_compile_many_outputs(self):
+
+        @mx.compile
+        def fun(arr):
+            arrs = [arr] * 64
+            first_arrs = None
+            for _ in range(6):
+                arrs = [x + y for x, y in zip(arrs[::2], arrs[1::2])]
+                if first_arrs is None:
+                    first_arrs = arrs
+            return arrs[0], first_arrs
+
+        out = fun(mx.array([1.0, 2.0]))
+        self.assertTrue(mx.allclose(out[0], mx.array([64.0, 128.0])))
+
     def test_shapeless_compile_matmul(self):
         a = mx.array([0.0, 1.0, 2.0])
         b = mx.array([0.0, 1.0, 2.0])
@@ -928,6 +953,7 @@ class TestCompile(mlx_tests.MLXTestCase):
         self.assertEqual(out[1].shape, (2, 2, 5))
 
     def test_leaks(self):
+        gc.collect()
         if mx.metal.is_available():
             mem_pre = mx.metal.get_active_memory()
         else:

python/tests/test_conv.py

@@ -341,7 +341,7 @@ class TestConv(mlx_tests.MLXTestCase):
                     atol, rtol = 1e-1, 1e-3
                 else:
                     atol, rtol = 1e-5, 1e-6
-                self.assertTrue(np.allclose(out_pt, out_mx, atol=atol, rtol=rtol))
+                self.assertTrue(np.allclose(out_pt, out_mx, atol=atol))
 
         for dtype in ("float32", "bfloat16"):
             for N, C, O in (
@@ -1042,6 +1042,14 @@ class TestConv(mlx_tests.MLXTestCase):
         self.assertTrue(mx.allclose(expected[0], grads[0]))
         self.assertTrue(mx.allclose(expected[1], grads[1]))
 
+    def test_repeated_conv(self):
+        x = mx.random.normal((1, 3, 3, 320))
+        w = mx.random.normal((320, 3, 3, 320))
+        for i in range(8):
+            y1 = mx.conv2d(x, w, (1, 1), (1, 1), (1, 1), 1)
+            y2 = mx.conv2d(x, w, (1, 1), (1, 1), (1, 1), 1)
+            self.assertTrue(mx.allclose(y1, y2))
+
 
 if __name__ == "__main__":
     unittest.main()

python/tests/test_double.py

@@ -173,6 +173,125 @@ class TestDouble(mlx_tests.MLXTestCase):
             mx.allclose(y, y_double.astype(mx.float32, mx.cpu), equal_nan=True)
         )
 
+    def test_type_promotion(self):
+        import mlx.core as mx
+
+        a = mx.array([4, 8], mx.float64)
+        b = mx.array([4, 8], mx.int32)
+
+        with mx.stream(mx.cpu):
+            c = a + b
+            self.assertEqual(c.dtype, mx.float64)
+
+    def test_lapack(self):
+        with mx.stream(mx.cpu):
+            # QRF
+            A = mx.array([[2.0, 3.0], [1.0, 2.0]], dtype=mx.float64)
+            Q, R = mx.linalg.qr(A)
+            out = Q @ R
+            self.assertTrue(mx.allclose(out, A))
+            out = Q.T @ Q
+            self.assertTrue(mx.allclose(out, mx.eye(2)))
+            self.assertTrue(mx.allclose(mx.tril(R, -1), mx.zeros_like(R)))
+            self.assertEqual(Q.dtype, mx.float64)
+            self.assertEqual(R.dtype, mx.float64)
+
+            # SVD
+            A = mx.array(
+                [[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]], dtype=mx.float64
+            )
+            U, S, Vt = mx.linalg.svd(A)
+            self.assertTrue(mx.allclose(U[:, : len(S)] @ mx.diag(S) @ Vt, A))
+
+            # Inverse
+            A = mx.array([[1, 2, 3], [6, -5, 4], [-9, 8, 7]], dtype=mx.float64)
+            A_inv = mx.linalg.inv(A)
+            self.assertTrue(mx.allclose(A @ A_inv, mx.eye(A.shape[0])))
+
+            # Tri inv
+            A = mx.array([[1, 0, 0], [6, -5, 0], [-9, 8, 7]], dtype=mx.float64)
+            B = mx.array([[7, 0, 0], [3, -2, 0], [1, 8, 3]], dtype=mx.float64)
+            AB = mx.stack([A, B])
+            invs = mx.linalg.tri_inv(AB, upper=False)
+            for M, M_inv in zip(AB, invs):
+                self.assertTrue(mx.allclose(M @ M_inv, mx.eye(M.shape[0])))
+
+            # Cholesky
+            sqrtA = mx.array(
+                [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]], dtype=mx.float64
+            )
+            A = sqrtA.T @ sqrtA / 81
+            L = mx.linalg.cholesky(A)
+            U = mx.linalg.cholesky(A, upper=True)
+            self.assertTrue(mx.allclose(L @ L.T, A))
+            self.assertTrue(mx.allclose(U.T @ U, A))
+
+            # Psueod inverse
+            A = mx.array([[1, 2, 3], [6, -5, 4], [-9, 8, 7]], dtype=mx.float64)
+            A_plus = mx.linalg.pinv(A)
+            self.assertTrue(mx.allclose(A @ A_plus @ A, A))
+
+            # Eigh
+            def check_eigs_and_vecs(A_np, kwargs={}):
+                A = mx.array(A_np, dtype=mx.float64)
+                eig_vals, eig_vecs = mx.linalg.eigh(A, **kwargs)
+                eig_vals_np, _ = np.linalg.eigh(A_np, **kwargs)
+                self.assertTrue(np.allclose(eig_vals, eig_vals_np))
+                self.assertTrue(
+                    mx.allclose(A @ eig_vecs, eig_vals[..., None, :] * eig_vecs)
+                )
+
+                eig_vals_only = mx.linalg.eigvalsh(A, **kwargs)
+                self.assertTrue(mx.allclose(eig_vals, eig_vals_only))
+
+            # Test a simple 2x2 symmetric matrix
+            A_np = np.array([[1.0, 2.0], [2.0, 4.0]], dtype=np.float64)
+            check_eigs_and_vecs(A_np)
+
+            # Test a larger random symmetric matrix
+            n = 5
+            np.random.seed(1)
+            A_np = np.random.randn(n, n).astype(np.float64)
+            A_np = (A_np + A_np.T) / 2
+            check_eigs_and_vecs(A_np)
+
+            # Test with upper triangle
+            check_eigs_and_vecs(A_np, {"UPLO": "U"})
+
+            # LU factorization
+            # Test 3x3 matrix
+            a = mx.array(
+                [[3.0, 1.0, 2.0], [1.0, 8.0, 6.0], [9.0, 2.0, 5.0]], dtype=mx.float64
+            )
+            P, L, U = mx.linalg.lu(a)
+            self.assertTrue(mx.allclose(L[P, :] @ U, a))
+
+            # Solve triangular
+            # Test lower triangular matrix
+            a = mx.array(
+                [[4.0, 0.0, 0.0], [2.0, 3.0, 0.0], [1.0, -2.0, 5.0]], dtype=mx.float64
+            )
+            b = mx.array([8.0, 14.0, 3.0], dtype=mx.float64)
+
+            result = mx.linalg.solve_triangular(a, b, upper=False)
+            expected = np.linalg.solve(np.array(a), np.array(b))
+            self.assertTrue(np.allclose(result, expected))
+
+            # Test upper triangular matrix
+            a = mx.array(
+                [[3.0, 2.0, 1.0], [0.0, 5.0, 4.0], [0.0, 0.0, 6.0]], dtype=mx.float64
+            )
+            b = mx.array([13.0, 33.0, 18.0], dtype=mx.float64)
+
+            result = mx.linalg.solve_triangular(a, b, upper=True)
+            expected = np.linalg.solve(np.array(a), np.array(b))
+            self.assertTrue(np.allclose(result, expected))
+
+    def test_conversion(self):
+        a = mx.array([1.0, 2.0], mx.float64)
+        b = np.array(a)
+        self.assertTrue(np.array_equal(a, b))
+
 
 if __name__ == "__main__":
     unittest.main()

python/tests/test_fast.py

@@ -158,7 +158,17 @@ class TestFast(mlx_tests.MLXTestCase):
         )
         self.assertLess(mx.abs(rx - rx_fast).max(), tolerances[mx.float32])
 
+        # Test raises with integer inputs
+        dims, _, base, scale, offset, traditional = defaults
+        x = (mx.random.uniform(shape=(2, T, dims)) * 10).astype(mx.int32)
+        with self.assertRaises(ValueError):
+            y = mx.fast.rope(
+                x, dims, traditional=traditional, base=base, scale=scale, offset=offset
+            )
+
     def test_rope_with_freqs(self):
+        mx.random.seed(0)
+
         # Check throws
         T = 4
         dims = 8
@@ -288,6 +298,9 @@ class TestFast(mlx_tests.MLXTestCase):
             rx = rms_norm(x, weight, eps)
             rx_fast = mx.fast.rms_norm(x, weight, eps)
             self.assertLess(mx.abs(rx - rx_fast).max(), tolerances[dtype])
+            rx = rms_norm(x, mx.ones_like(weight), eps)
+            rx_fast = mx.fast.rms_norm(x, None, eps)
+            self.assertLess(mx.abs(rx - rx_fast).max(), tolerances[dtype])
 
         for eps in epss:
             dtype, _, dims = defaults
@@ -296,6 +309,9 @@ class TestFast(mlx_tests.MLXTestCase):
             rx = rms_norm(x, weight, eps)
             rx_fast = mx.fast.rms_norm(x, weight, eps)
             self.assertLess(mx.abs(rx - rx_fast).max(), tolerances[dtype])
+            rx = rms_norm(x, mx.ones_like(weight), eps)
+            rx_fast = mx.fast.rms_norm(x, None, eps)
+            self.assertLess(mx.abs(rx - rx_fast).max(), tolerances[dtype])
 
         for dims in dimss:
             dtype, eps, _ = defaults
@@ -304,6 +320,9 @@ class TestFast(mlx_tests.MLXTestCase):
             rx = rms_norm(x, weight, eps)
             rx_fast = mx.fast.rms_norm(x, weight, eps)
             self.assertLess(mx.abs(rx - rx_fast).max(), tolerances[dtype])
+            rx = rms_norm(x, mx.ones_like(weight), eps)
+            rx_fast = mx.fast.rms_norm(x, None, eps)
+            self.assertLess(mx.abs(rx - rx_fast).max(), tolerances[dtype])
 
         # Test > 4096
         dims, dtype, eps = 4099, mx.float32, 1e-5
@@ -323,6 +342,8 @@ class TestFast(mlx_tests.MLXTestCase):
         eps = 1e-5
         f1 = lambda x, w, y: (rms_norm(x, w, eps) * y).sum()
         f2 = lambda x, w, y: (mx.fast.rms_norm(x, w, eps) * y).sum()
+        f3 = lambda x, y: (rms_norm(x, mx.ones((x.shape[-1],)), eps) * y).sum()
+        f4 = lambda x, y: (mx.fast.rms_norm(x, None, eps) * y).sum()
 
         x = mx.random.uniform(shape=(8, 100, D))
         w = mx.random.uniform(shape=(D,))
@@ -331,6 +352,9 @@ class TestFast(mlx_tests.MLXTestCase):
         gx2, gw2 = mx.grad(f2, argnums=(0, 1))(x, w, y)
         self.assertLess(mx.abs(gx1 - gx2).max(), 1e-5)
         self.assertLess(mx.abs(gw1 - gw2).max() / mx.abs(gw1).mean(), 1e-5)
+        gx1 = mx.grad(f3, argnums=(0,))(x, y)
+        gx2 = mx.grad(f4, argnums=(0,))(x, y)
+        self.assertLess(mx.abs(gx1 - gx2).max(), 1e-5)
 
         D = 8192
         x = mx.random.uniform(shape=(2, 2, D))
@@ -340,6 +364,9 @@ class TestFast(mlx_tests.MLXTestCase):
         gx2, gw2 = mx.grad(f2, argnums=(0, 1))(x, w, y)
         self.assertLess(mx.abs(gx1 - gx2).max(), 1e-5)
         self.assertLess(mx.abs(gw1 - gw2).max() / mx.abs(gw1).mean(), 1e-5)
+        gx1 = mx.grad(f3, argnums=(0,))(x, y)
+        gx2 = mx.grad(f4, argnums=(0,))(x, y)
+        self.assertLess(mx.abs(gx1 - gx2).max(), 1e-5)
 
         def gf(f):
             def inner(x, w, y):

python/tests/test_fast_sdpa.py

@@ -262,6 +262,61 @@ class TestFastSDPA(mlx_tests.MLXTestCase):
             )
             self.assertTrue(mx.allclose(ref, out, atol=1e-4, rtol=1e-4))
 
+    def test_fast_sdpa_few_query(self):
+        D = 64
+        L = 43
+        Lq = 4
+        Nq = 8
+        Nkv = 1
+        scale = 1.0
+        mx.random.seed(0)
+        q = 5e-1 * mx.random.normal(shape=(1, Lq, Nq, D))
+        q = q.swapaxes(1, 2)
+        k = 5e-1 * mx.random.normal(shape=(1, Nkv, L, D))
+        v = 5e-1 * mx.random.normal(shape=(1, Nkv, L, D))
+
+        masks = [
+            mx.array(True),
+            mx.array([True] * (L - 10) + [False] * 10),
+            mx.random.uniform(shape=(Nq, 1, L)) > 0.2,
+            mx.random.uniform(shape=(L, 1, Nq)).T > 0.2,
+        ]
+        for m in masks:
+            ref = mlx_primitives_sdpa(q, k, v, scale, mask=m)
+            out = mx.fast.scaled_dot_product_attention(
+                q,
+                k,
+                v,
+                scale=scale,
+                mask=m,
+            )
+            self.assertTrue(mx.allclose(ref, out, atol=1e-4, rtol=1e-4))
+
+        return
+        L = 4096
+        scale = 1.0
+        mx.random.seed(0)
+        q = 5e-1 * mx.random.normal(shape=(1, Nq, Lq, D))
+        k = 5e-1 * mx.random.normal(shape=(1, Nkv, L, D))
+        v = 5e-1 * mx.random.normal(shape=(1, Nkv, L, D))
+
+        masks = [
+            mx.array(True),
+            mx.array([True] * (L - 10) + [False] * 10),
+            mx.random.uniform(shape=(Nq, 1, L)) > 0.2,
+            mx.random.uniform(shape=(L, 1, Nq)).T > 0.2,
+        ]
+        for m in masks:
+            ref = mlx_primitives_sdpa(q, k, v, scale, mask=m)
+            out = mx.fast.scaled_dot_product_attention(
+                q,
+                k,
+                v,
+                scale=scale,
+                mask=m,
+            )
+            self.assertTrue(mx.allclose(ref, out, atol=1e-4, rtol=1e-4))
+
     @unittest.skip("Different head and value dims is not enabled")
     def test_fast_sdpa_vector_value_dims(self):
         D = 192

python/tests/test_linalg.py

@@ -12,11 +12,11 @@ import numpy as np
 class TestLinalg(mlx_tests.MLXTestCase):
     def test_norm(self):
         vector_ords = [None, 0.5, 0, 1, 2, 3, -1, float("inf"), -float("inf")]
-        matrix_ords = [None, "fro", -1, 1, float("inf"), -float("inf")]
+        matrix_ords = [None, "fro", "nuc", -1, 1, -2, 2, float("inf"), -float("inf")]
 
         for shape in [(3,), (2, 3), (2, 3, 3)]:
-            x_mx = mx.arange(1, math.prod(shape) + 1).reshape(shape)
-            x_np = np.arange(1, math.prod(shape) + 1).reshape(shape)
+            x_mx = mx.arange(1, math.prod(shape) + 1, dtype=mx.float32).reshape(shape)
+            x_np = np.arange(1, math.prod(shape) + 1, dtype=np.float32).reshape(shape)
             # Test when at least one axis is provided
             for num_axes in range(1, len(shape)):
                 if num_axes == 1:
@@ -26,11 +26,14 @@ class TestLinalg(mlx_tests.MLXTestCase):
                 for axis in itertools.combinations(range(len(shape)), num_axes):
                     for keepdims in [True, False]:
                         for o in ords:
+                            stream = (
+                                mx.cpu if o in ["nuc", -2, 2] else mx.default_device()
+                            )
                             out_np = np.linalg.norm(
                                 x_np, ord=o, axis=axis, keepdims=keepdims
                             )
                             out_mx = mx.linalg.norm(
-                                x_mx, ord=o, axis=axis, keepdims=keepdims
+                                x_mx, ord=o, axis=axis, keepdims=keepdims, stream=stream
                             )
                             with self.subTest(
                                 shape=shape, ord=o, axis=axis, keepdims=keepdims
@@ -133,20 +136,38 @@ class TestLinalg(mlx_tests.MLXTestCase):
 
     def test_svd_decomposition(self):
         A = mx.array([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]], dtype=mx.float32)
-        U, S, Vt = mx.linalg.svd(A, stream=mx.cpu)
+        U, S, Vt = mx.linalg.svd(A, compute_uv=True, stream=mx.cpu)
         self.assertTrue(
             mx.allclose(U[:, : len(S)] @ mx.diag(S) @ Vt, A, rtol=1e-5, atol=1e-7)
         )
 
+        S = mx.linalg.svd(A, compute_uv=False, stream=mx.cpu)
+        self.assertTrue(
+            mx.allclose(
+                mx.linalg.norm(S), mx.linalg.norm(A, ord="fro"), rtol=1e-5, atol=1e-7
+            )
+        )
+
         # Multiple matrices
         B = A + 10.0
         AB = mx.stack([A, B])
-        Us, Ss, Vts = mx.linalg.svd(AB, stream=mx.cpu)
+        Us, Ss, Vts = mx.linalg.svd(AB, compute_uv=True, stream=mx.cpu)
         for M, U, S, Vt in zip([A, B], Us, Ss, Vts):
             self.assertTrue(
                 mx.allclose(U[:, : len(S)] @ mx.diag(S) @ Vt, M, rtol=1e-5, atol=1e-7)
             )
 
+        Ss = mx.linalg.svd(AB, compute_uv=False, stream=mx.cpu)
+        for M, S in zip([A, B], Ss):
+            self.assertTrue(
+                mx.allclose(
+                    mx.linalg.norm(S),
+                    mx.linalg.norm(M, ord="fro"),
+                    rtol=1e-5,
+                    atol=1e-7,
+                )
+            )
+
     def test_inverse(self):
         A = mx.array([[1, 2, 3], [6, -5, 4], [-9, 8, 7]], dtype=mx.float32)
         A_inv = mx.linalg.inv(A, stream=mx.cpu)
@@ -175,6 +196,13 @@ class TestLinalg(mlx_tests.MLXTestCase):
                     mx.allclose(M @ M_inv, mx.eye(M.shape[0]), rtol=0, atol=1e-5)
                 )
 
+        # Ensure that tri_inv will 0-out the supposedly 0 triangle
+        x = mx.random.normal((2, 8, 8))
+        y1 = mx.linalg.tri_inv(x, upper=True, stream=mx.cpu)
+        y2 = mx.linalg.tri_inv(x, upper=False, stream=mx.cpu)
+        self.assertTrue(mx.all(y1 == mx.triu(y1)))
+        self.assertTrue(mx.all(y2 == mx.tril(y2)))
+
     def test_cholesky(self):
         sqrtA = mx.array(
             [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]], dtype=mx.float32
@@ -351,6 +379,15 @@ class TestLinalg(mlx_tests.MLXTestCase):
         L = mx.take_along_axis(L, P[..., None], axis=-2)
         self.assertTrue(mx.allclose(L @ U, a))
 
+        # Test non-square matrix
+        a = mx.array([[3.0, 1.0, 2.0], [1.0, 8.0, 6.0]])
+        P, L, U = mx.linalg.lu(a, stream=mx.cpu)
+        self.assertTrue(mx.allclose(L[P, :] @ U, a))
+
+        a = mx.array([[3.0, 1.0], [1.0, 8.0], [9.0, 2.0]])
+        P, L, U = mx.linalg.lu(a, stream=mx.cpu)
+        self.assertTrue(mx.allclose(L[P, :] @ U, a))
+
     def test_lu_factor(self):
         mx.random.seed(7)
 

python/tests/test_load.py

@@ -385,6 +385,7 @@ class TestLoad(mlx_tests.MLXTestCase):
         mx.eval(x)
         save_file = os.path.join(self.test_dir, "donation.npy")
         mx.save(save_file, x)
+        mx.synchronize(mx.default_stream(mx.default_device()))
 
         mx.metal.reset_peak_memory()
         scale = mx.array(2.0)

python/tests/test_ops.py

@@ -898,6 +898,10 @@ class TestOps(mlx_tests.MLXTestCase):
         ).astype(np.float32)
         self.assertTrue(np.allclose(mx.erfinv(x), expected, equal_nan=True))
 
+        result = mx.erfinv(mx.array([0.9999999403953552] * 8))
+        expected = mx.array([3.8325066566467285] * 8)
+        self.assertTrue(mx.allclose(result, expected))
+
     def test_sin(self):
         a = mx.array(
             [0, math.pi / 4, math.pi / 2, math.pi, 3 * math.pi / 4, 2 * math.pi]
@@ -1890,6 +1894,22 @@ class TestOps(mlx_tests.MLXTestCase):
             expected = mx.repeat(expected[:, None], 2, axis=1)
             self.assertTrue(mx.array_equal(expected, out))
 
+        # Test donation
+        def fn(its):
+            x = mx.ones((32,))
+            for _ in range(its):
+                x = mx.cumsum(x)
+            return x
+
+        mx.synchronize(mx.default_stream(mx.default_device()))
+        mx.eval(fn(2))
+        mx.synchronize(mx.default_stream(mx.default_device()))
+        mem2 = mx.metal.get_peak_memory()
+        mx.eval(fn(4))
+        mx.synchronize(mx.default_stream(mx.default_device()))
+        mem4 = mx.metal.get_peak_memory()
+        self.assertEqual(mem2, mem4)
+
     def test_squeeze_expand(self):
         a = mx.zeros((2, 1, 2, 1))
         self.assertEqual(mx.squeeze(a).shape, (2, 2))
@@ -2189,6 +2209,14 @@ class TestOps(mlx_tests.MLXTestCase):
         expected = mx.array(np.linspace(1, 10, 1))
         self.assertEqualArray(d, expected)
 
+        # Ensure that the start and stop are always the ones provided
+        ranges = mx.random.normal((16, 2)).tolist()
+        nums = (2 + mx.random.uniform(shape=(16,)) * 10).astype(mx.uint32).tolist()
+        for (a, b), n in zip(ranges, nums):
+            d = mx.linspace(a, b, n).tolist()
+            self.assertEqual(d[0], a)
+            self.assertEqual(d[-1], b)
+
     def test_repeat(self):
         # Setup data for the tests
         data = mx.array([[[13, 3], [16, 6]], [[14, 4], [15, 5]], [[11, 1], [12, 2]]])
@@ -2834,6 +2862,11 @@ class TestOps(mlx_tests.MLXTestCase):
         b[::2] = 0
         self.assertTrue(mx.array_equal(b, mx.array([0, 3, 0, 1])))
 
+    def test_slice_with_negative_stride(self):
+        a = mx.random.uniform(shape=(128, 4))
+        out = a[::-1]
+        self.assertTrue(mx.array_equal(out[-1, :], a[0, :]))
+
 
 if __name__ == "__main__":
     unittest.main()

python/tests/test_optimizers.py

@@ -39,6 +39,7 @@ def tree_equal(fn, *args):
 
 
 optimizers_dict = get_all_optimizers()
+del optimizers_dict["MultiOptimizer"]
 
 
 class TestOptimizers(mlx_tests.MLXTestCase):
@@ -500,6 +501,30 @@ class TestSchedulers(unittest.TestCase):
         grads = model.trainable_parameters()
         optimizer.update(model, grads)
 
+    def test_multi_optimizer(self):
+        class Model(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.l1 = nn.Linear(2, 2)
+                self.drop = nn.Dropout(p=0.5)
+                self.l2 = nn.Linear(2, 2)
+                self.vals = [nn.Linear(2, 2), nn.ReLU(), nn.ReLU()]
+
+        model = Model()
+        optimizer = opt.MultiOptimizer(
+            [opt.Adam(learning_rate=0.001), opt.SGD(learning_rate=0.1)],
+            [lambda name, weight: weight.ndim > 1],
+        )
+        optimizer.init(model.trainable_parameters())
+
+        self.assertEqual(len(optimizer.state["states"]), 2)
+
+        adam_states = tree_flatten(optimizer.state["states"][0])
+        sgd_states = tree_flatten(optimizer.state["states"][1])
+        self.assertEqual((len(sgd_states) - 2) * 2, len(adam_states) - 2)
+        self.assertFalse(any("bias" in k for k, v in adam_states))
+        self.assertFalse(any("weight" in k for k, v in sgd_states))
+
 
 if __name__ == "__main__":
     unittest.main()

python/tests/test_tree.py

@@ -3,6 +3,7 @@
 import unittest
 
 import mlx.core as mx
+import mlx.nn as nn
 import mlx.utils
 import mlx_tests
 
@@ -22,6 +23,29 @@ class TestTreeUtils(mlx_tests.MLXTestCase):
         self.assertEqual(list(zip(*flat_tree))[1], vals)
         self.assertEqual(mlx.utils.tree_unflatten(flat_tree), tree)
 
+    def test_merge(self):
+        t1 = {"a": 0}
+        t2 = {"b": 1}
+        t = mlx.utils.tree_merge(t1, t2)
+        self.assertEqual({"a": 0, "b": 1}, t)
+        with self.assertRaises(ValueError):
+            mlx.utils.tree_merge(t1, t1)
+        with self.assertRaises(ValueError):
+            mlx.utils.tree_merge(t, t1)
+
+        mod1 = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
+        mod2 = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
+        mod = nn.Sequential(mod1, mod2)
+
+        params1 = {"layers": [mod1.parameters()]}
+        params2 = {"layers": [None, mod2.parameters()]}
+        params = mlx.utils.tree_merge(params1, params2)
+        for (k1, v1), (k2, v2) in zip(
+            mlx.utils.tree_flatten(params), mlx.utils.tree_flatten(mod.parameters())
+        ):
+            self.assertEqual(k1, k2)
+            self.assertTrue(mx.array_equal(v1, v2))
+
 
 if __name__ == "__main__":
     unittest.main()

python/tests/test_vmap.py

@@ -316,35 +316,59 @@ class TestVmap(mlx_tests.MLXTestCase):
     def test_vmap_svd(self):
         a = mx.random.uniform(shape=(3, 4, 2))
 
-        cpu_svd = lambda x: mx.linalg.svd(x, stream=mx.cpu)
+        cpu_svd_full = lambda x: mx.linalg.svd(x, compute_uv=True, stream=mx.cpu)
+        cpu_svd_singular = lambda x: mx.linalg.svd(x, compute_uv=False, stream=mx.cpu)
 
         # Vmap over the first axis (this is already supported natively by the primitive).
-        Us, Ss, Vts = mx.vmap(cpu_svd, in_axes=(0,))(a)
+        Us, Ss, Vts = mx.vmap(cpu_svd_full, in_axes=(0,))(a)
         self.assertEqual(Us.shape, (a.shape[0], a.shape[1], a.shape[1]))
         self.assertEqual(Ss.shape, (a.shape[0], a.shape[2]))
         self.assertEqual(Vts.shape, (a.shape[0], a.shape[2], a.shape[2]))
 
+        Sv = mx.vmap(cpu_svd_singular, in_axes=(0,))(a)
+        self.assertEqual(Sv.shape, (a.shape[0], a.shape[2]))
+
         for i in range(a.shape[0]):
             M = a[i]
             U, S, Vt = Us[i], Ss[i], Vts[i]
             self.assertTrue(
                 mx.allclose(U[:, : len(S)] @ mx.diag(S) @ Vt, M, rtol=1e-5, atol=1e-7)
             )
+            self.assertTrue(
+                mx.allclose(
+                    mx.linalg.norm(Sv[i]),
+                    mx.linalg.norm(M, ord="fro"),
+                    rtol=1e-5,
+                    atol=1e-7,
+                )
+            )
 
         # Vmap over the second axis.
-        Us, Ss, Vts = mx.vmap(cpu_svd, in_axes=(1,))(a)
+        Us, Ss, Vts = mx.vmap(cpu_svd_full, in_axes=(1,))(a)
         self.assertEqual(Us.shape, (a.shape[1], a.shape[0], a.shape[0]))
         self.assertEqual(Ss.shape, (a.shape[1], a.shape[2]))
         self.assertEqual(Vts.shape, (a.shape[1], a.shape[2], a.shape[2]))
 
+        Sv = mx.vmap(cpu_svd_singular, in_axes=(1,))(a)
+        self.assertEqual(Sv.shape, (a.shape[1], a.shape[2]))
+
         for i in range(a.shape[1]):
             M = a[:, i, :]
             U, S, Vt = Us[i], Ss[i], Vts[i]
             self.assertTrue(
                 mx.allclose(U[:, : len(S)] @ mx.diag(S) @ Vt, M, rtol=1e-5, atol=1e-7)
             )
+            self.assertTrue(
+                mx.allclose(
+                    mx.linalg.norm(Sv[i]),
+                    mx.linalg.norm(M, ord="fro"),
+                    rtol=1e-5,
+                    atol=1e-7,
+                )
+            )
 
     def test_vmap_inverse(self):
+        mx.random.seed(42)
         a = mx.random.uniform(shape=(3, 4, 4))
 
         cpu_inv = lambda x: mx.linalg.inv(x, stream=mx.cpu)

setup.py

@@ -173,7 +173,7 @@ if __name__ == "__main__":
 
     setup(
         name="mlx",
-        version=get_version("0.23.0"),
+        version=get_version("0.23.2"),
         author="MLX Contributors",
         author_email="mlx@group.apple.com",
         description="A framework for machine learning on Apple silicon.",
@@ -194,7 +194,12 @@ if __name__ == "__main__":
                 "typing_extensions",
             ],
         },
-        entry_points={"console_scripts": ["mlx.launch = mlx.distributed_run:main"]},
+        entry_points={
+            "console_scripts": [
+                "mlx.launch = mlx.distributed_run:main",
+                "mlx.distributed_config = mlx.distributed_run:distributed_config",
+            ]
+        },
         ext_modules=[CMakeExtension("mlx.core")],
         cmdclass={"build_ext": CMakeBuild, "generate_stubs": GenerateStubs},
         zip_safe=False,

tests/linalg_tests.cpp

@@ -100,7 +100,7 @@ TEST_CASE("[mlx.core.linalg.norm] double ord") {
       norm(x, -std::numeric_limits<double>::infinity()).item<float>(),
       doctest::Approx(expected));
 
-  x = reshape(arange(9), {3, 3});
+  x = reshape(arange(9, float32), {3, 3});
 
   CHECK(allclose(
             norm(x, 2.0, 0, false),
@@ -129,10 +129,34 @@ TEST_CASE("[mlx.core.linalg.norm] double ord") {
   CHECK_EQ(
       norm(x, -1.0, std::vector<int>{1, 0}).item<float>(),
       doctest::Approx(3.0));
+  CHECK_EQ(
+      norm(x, 2.0, std::vector<int>{0, 1}, false, Device::cpu).item<float>(),
+      doctest::Approx(14.226707));
+  CHECK_EQ(
+      norm(x, 2.0, std::vector<int>{1, 0}, false, Device::cpu).item<float>(),
+      doctest::Approx(14.226707));
+  CHECK_EQ(
+      norm(x, -2.0, std::vector<int>{0, 1}, false, Device::cpu).item<float>(),
+      doctest::Approx(0.0));
+  CHECK_EQ(
+      norm(x, -2.0, std::vector<int>{1, 0}, false, Device::cpu).item<float>(),
+      doctest::Approx(0.0));
   CHECK_EQ(norm(x, 1.0, std::vector<int>{0, 1}, true).shape(), Shape{1, 1});
   CHECK_EQ(norm(x, 1.0, std::vector<int>{1, 0}, true).shape(), Shape{1, 1});
   CHECK_EQ(norm(x, -1.0, std::vector<int>{0, 1}, true).shape(), Shape{1, 1});
   CHECK_EQ(norm(x, -1.0, std::vector<int>{1, 0}, true).shape(), Shape{1, 1});
+  CHECK_EQ(
+      norm(x, 2.0, std::vector<int>{0, 1}, true, Device::cpu).shape(),
+      Shape{1, 1});
+  CHECK_EQ(
+      norm(x, 2.0, std::vector<int>{1, 0}, true, Device::cpu).shape(),
+      Shape{1, 1});
+  CHECK_EQ(
+      norm(x, -2.0, std::vector<int>{0, 1}, true, Device::cpu).shape(),
+      Shape{1, 1});
+  CHECK_EQ(
+      norm(x, -2.0, std::vector<int>{1, 0}, true, Device::cpu).shape(),
+      Shape{1, 1});
 
   CHECK_EQ(
       norm(x, -1.0, std::vector<int>{-2, -1}, false).item<float>(),
@@ -140,8 +164,14 @@ TEST_CASE("[mlx.core.linalg.norm] double ord") {
   CHECK_EQ(
       norm(x, 1.0, std::vector<int>{-2, -1}, false).item<float>(),
       doctest::Approx(15.0));
+  CHECK_EQ(
+      norm(x, -2.0, std::vector<int>{-2, -1}, false, Device::cpu).item<float>(),
+      doctest::Approx(0.0));
+  CHECK_EQ(
+      norm(x, 2.0, std::vector<int>{-2, -1}, false, Device::cpu).item<float>(),
+      doctest::Approx(14.226707));
 
-  x = reshape(arange(18), {2, 3, 3});
+  x = reshape(arange(18, float32), {2, 3, 3});
   CHECK_THROWS(norm(x, 2.0, std::vector{0, 1, 2}));
   CHECK(allclose(
             norm(x, 3.0, 0),
@@ -199,13 +229,31 @@ TEST_CASE("[mlx.core.linalg.norm] double ord") {
             .item<bool>());
   CHECK(allclose(norm(x, -1.0, std::vector<int>{1, 2}), array({9, 36}))
             .item<bool>());
+  CHECK(allclose(
+            norm(x, 2.0, std::vector<int>{0, 1}, false, Device::cpu),
+            array({22.045408, 24.155825, 26.318918}))
+            .item<bool>());
+  CHECK(allclose(
+            norm(x, 2.0, std::vector<int>{1, 2}, false, Device::cpu),
+            array({14.226707, 39.759212}))
+            .item<bool>());
+  CHECK(allclose(
+            norm(x, -2.0, std::vector<int>{0, 1}, false, Device::cpu),
+            array({3, 2.7378995, 2.5128777}))
+            .item<bool>());
+  CHECK(allclose(
+            norm(x, -2.0, std::vector<int>{1, 2}, false, Device::cpu),
+            array({4.979028e-16, 7.009628e-16}),
+            /* rtol = */ 1e-5,
+            /* atol = */ 1e-6)
+            .item<bool>());
 }
 
 TEST_CASE("[mlx.core.linalg.norm] string ord") {
   array x({1, 2, 3});
   CHECK_THROWS(norm(x, "fro"));
 
-  x = reshape(arange(9), {3, 3});
+  x = reshape(arange(9, float32), {3, 3});
   CHECK_THROWS(norm(x, "bad ord"));
 
   CHECK_EQ(
@@ -214,8 +262,11 @@ TEST_CASE("[mlx.core.linalg.norm] string ord") {
   CHECK_EQ(
       norm(x, "fro", std::vector<int>{0, 1}).item<float>(),
       doctest::Approx(14.2828568570857));
+  CHECK_EQ(
+      norm(x, "nuc", std::vector<int>{0, 1}, false, Device::cpu).item<float>(),
+      doctest::Approx(15.491934));
 
-  x = reshape(arange(18), {2, 3, 3});
+  x = reshape(arange(18, float32), {2, 3, 3});
   CHECK(allclose(
             norm(x, "fro", std::vector<int>{0, 1}),
             array({22.24859546, 24.31049156, 26.43860813}))
@@ -240,6 +291,18 @@ TEST_CASE("[mlx.core.linalg.norm] string ord") {
             norm(x, "f", std::vector<int>{2, 1}),
             array({14.28285686, 39.7617907}))
             .item<bool>());
+  CHECK(allclose(
+            norm(x, "nuc", std::vector<int>{0, 1}, false, Device::cpu),
+            array({25.045408, 26.893724, 28.831797}))
+            .item<bool>());
+  CHECK(allclose(
+            norm(x, "nuc", std::vector<int>{1, 2}, false, Device::cpu),
+            array({15.491934, 40.211937}))
+            .item<bool>());
+  CHECK(allclose(
+            norm(x, "nuc", std::vector<int>{-2, -1}, false, Device::cpu),
+            array({15.491934, 40.211937}))
+            .item<bool>());
 }
 
 TEST_CASE("test QR factorization") {
@@ -271,7 +334,7 @@ TEST_CASE("test SVD factorization") {
 
   const auto prng_key = random::key(42);
   const auto A = mlx::core::random::normal({5, 4}, prng_key);
-  const auto outs = linalg::svd(A, Device::cpu);
+  const auto outs = linalg::svd(A, true, Device::cpu);
   CHECK_EQ(outs.size(), 3);
 
   const auto& U = outs[0];
@@ -291,6 +354,15 @@ TEST_CASE("test SVD factorization") {
   CHECK_EQ(U.dtype(), float32);
   CHECK_EQ(S.dtype(), float32);
   CHECK_EQ(Vt.dtype(), float32);
+
+  // Test singular values
+  const auto& outs_sv = linalg::svd(A, false, Device::cpu);
+  const auto SV = outs_sv[0];
+
+  CHECK_EQ(SV.shape(), Shape{4});
+  CHECK_EQ(SV.dtype(), float32);
+
+  CHECK(allclose(norm(SV), norm(A, "fro")).item<bool>());
 }
 
 TEST_CASE("test matrix inversion") {

tests/vmap_tests.cpp

@@ -466,15 +466,19 @@ TEST_CASE("test vmap scatter") {
 }
 
 TEST_CASE("test vmap SVD") {
-  auto fun = [](std::vector<array> inputs) {
-    return linalg::svd(inputs.at(0), Device::cpu);
+  auto svd_full = [](std::vector<array> inputs) {
+    return linalg::svd(inputs.at(0), true, Device::cpu);
+  };
+
+  auto svd_singular = [](std::vector<array> inputs) {
+    return linalg::svd(inputs.at(0), false, Device::cpu);
   };
 
   auto a = astype(reshape(arange(24), {3, 4, 2}), float32);
 
   // vmap over the second axis.
   {
-    auto out = vmap(fun, /* in_axes = */ {1})({a});
+    auto out = vmap(svd_full, /* in_axes = */ {1})({a});
     const auto& U = out.at(0);
     const auto& S = out.at(1);
     const auto& Vt = out.at(2);
@@ -486,7 +490,7 @@ TEST_CASE("test vmap SVD") {
 
   // vmap over the third axis.
   {
-    auto out = vmap(fun, /* in_axes = */ {2})({a});
+    auto out = vmap(svd_full, /* in_axes = */ {2})({a});
     const auto& U = out.at(0);
     const auto& S = out.at(1);
     const auto& Vt = out.at(2);
@@ -495,6 +499,21 @@ TEST_CASE("test vmap SVD") {
     CHECK_EQ(S.shape(), Shape{a.shape(2), a.shape(0)});
     CHECK_EQ(Vt.shape(), Shape{a.shape(2), a.shape(1), a.shape(1)});
   }
+
+  // test singular values
+  {
+    auto out = vmap(svd_singular, /* in_axes = */ {1})({a});
+    const auto& S = out.at(0);
+
+    CHECK_EQ(S.shape(), Shape{a.shape(1), a.shape(2)});
+  }
+
+  {
+    auto out = vmap(svd_singular, /* in_axes = */ {2})({a});
+    const auto& S = out.at(0);
+
+    CHECK_EQ(S.shape(), Shape{a.shape(2), a.shape(0)});
+  }
 }
 
 TEST_CASE("test vmap dynamic slices") {