Merge branch 'ml-explore:main' into rocm-support

Author: Geramy

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

@@ -45,15 +45,17 @@ runs:
         cd build
         cmake ..
         make -j $(sysctl -n hw.ncpu)
-    
+
     - name: Run CPP tests
       shell: bash -l {0}
       env:
         DEVICE: gpu
         METAL_DEVICE_WRAPPER_TYPE: 1
         METAL_DEBUG_ERROR_MODE: 0
-      run: ./build/tests/tests
-    
+      run: |
+        ./build/tests/tests
+        ./build/tests/test_teardown
+
     - name: Build small binary with JIT
       shell: bash -l {0}
       run: |

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

@@ -54,6 +54,12 @@ runs:
         echo PYTHONPATH=`python -c 'import sys; print(sys.path[-1])'` >> $GITHUB_ENV
         echo "::endgroup::"
 
+    - name: Set swap space
+      if: ${{ startsWith(inputs.toolkit, 'cuda') && runner.arch == 'arm64' }}
+      uses: pierotofy/set-swap-space@fc79b3f67fa8a838184ce84a674ca12238d2c761
+      with:
+        swap-size-gb: 16
+
     - name: Install CUDA toolkit
       if: ${{ startsWith(inputs.toolkit, 'cuda') }}
       shell: bash

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

@@ -17,4 +17,5 @@ runs:
       run: |
         echo "::group::CPP tests - CPU"
         ./build/tests.exe -tce="*gguf*,test random uniform"
+        ./build/test_teardown.exe
         echo "::endgroup::"

.github/workflows/documentation.yml

@@ -25,4 +25,4 @@ jobs:
     steps:
       - name: Deploy to GitHub Pages
         id: deployment
-        uses: actions/deploy-pages@v4
+        uses: actions/deploy-pages@v5

.github/workflows/release.yml

@@ -41,7 +41,7 @@ jobs:
     steps:
       - name: Deploy to GitHub Pages
         id: deployment
-        uses: actions/deploy-pages@v4
+        uses: actions/deploy-pages@v5
 
   build_linux_release:
     if: github.repository == 'ml-explore/mlx'

docs/src/dev/extensions.rst

@@ -404,7 +404,7 @@ below.
         auto kernel = d.get_kernel(kname, lib);
 
         // Prepare to encode kernel
-        auto& compute_encoder = d.get_command_encoder(s.index);
+        auto& compute_encoder = mx::metal::get_command_encoder(s);
         compute_encoder.set_compute_pipeline_state(kernel);
 
         // Kernel parameters are registered with buffer indices corresponding to
@@ -448,7 +448,7 @@ We can now call the :meth:`axpby` operation on both the CPU and the GPU!
 
 A few things to note about MLX and Metal before moving on. MLX keeps track of
 the active ``command_buffer`` and the ``MTLCommandBuffer`` to which it is
-associated. We rely on :meth:`d.get_command_encoder` to give us the active
+associated. We rely on :meth:`metal::get_command_encoder` to give us the active
 metal compute command encoder instead of building a new one and calling
 :meth:`compute_encoder->end_encoding` at the end. MLX adds kernels (compute
 pipelines) to the active command buffer until some specified limit is hit or

docs/src/index.rst

@@ -32,7 +32,7 @@ are the CPU and GPU.
    install
 
 .. toctree::
-   :caption: Usage 
+   :caption: Usage
    :maxdepth: 1
 
    usage/quick_start
@@ -78,6 +78,7 @@ are the CPU and GPU.
    python/optimizers
    python/distributed
    python/tree_utils
+   python/printoptions
 
 .. toctree::
    :caption: C++ API Reference

docs/src/python/fft.rst

@@ -20,5 +20,7 @@ FFT
   irfft2
   rfftn
   irfftn
+  fftfreq
+  rfftfreq
   fftshift
   ifftshift

docs/src/python/printoptions.rst

@@ -0,0 +1,12 @@
+Print Options
+===============
+
+.. currentmodule:: mlx.core
+
+.. autosummary::
+  :toctree: _autosummary
+
+  PrintOptions
+  set_printoptions
+  printoptions
+  get_printoptions

examples/extensions/axpby/axpby.cpp

@@ -192,7 +192,7 @@ void Axpby::eval_gpu(
   auto kernel = d.get_kernel(kname, lib);
 
   // Prepare to encode kernel
-  auto& compute_encoder = d.get_command_encoder(s.index);
+  auto& compute_encoder = mx::metal::get_command_encoder(s);
   compute_encoder.set_compute_pipeline_state(kernel);
 
   // Kernel parameters are registered with buffer indices corresponding to