diff --git a/.github/actions/build-cuda-release/action.yml b/.github/actions/build-cuda-release/action.yml
index 554e0869cf..71cbdb3f45 100644
--- a/.github/actions/build-cuda-release/action.yml
+++ b/.github/actions/build-cuda-release/action.yml
@@ -15,10 +15,12 @@ runs:
         MLX_BUILD_STAGE: 2
         CMAKE_ARGS: -DMLX_BUILD_CUDA=ON -DCMAKE_CUDA_COMPILER=${{ inputs.nvcc-location }}
       run: |
+        echo "::group::Build packages"
         pip install auditwheel build patchelf setuptools
         python setup.py clean --all
         python -m build -w
 
         if [ -f "python/scripts/repair_cuda.sh" ]; then
           bash python/scripts/repair_cuda.sh
-        fi
\ No newline at end of file
+        fi
+        echo "::endgroup::"
diff --git a/.github/actions/build-cuda/action.yml b/.github/actions/build-cuda/action.yml
index d6068247a4..b454bfa451 100644
--- a/.github/actions/build-cuda/action.yml
+++ b/.github/actions/build-cuda/action.yml
@@ -24,7 +24,10 @@ runs:
       env:
         DEBUG: 1
         CMAKE_ARGS: -DMLX_BUILD_CUDA=ON -DCMAKE_COMPILE_WARNING_AS_ERROR=ON -DCMAKE_CUDA_COMPILER=${{ inputs.nvcc-location }}
-      run: pip install -e ".[dev]" -v
+      run: |
+        echo "::group::Install Python Package"
+        pip install -e ".[dev]" -v
+        echo "::endgroup::"
 
     - name: Run Python tests - CPU
       if: inputs.run-tests == 'true'
@@ -32,7 +35,10 @@ runs:
       env:
         LOW_MEMORY: 1
         DEVICE: cpu
-      run: python -m unittest discover python/tests -v
+      run: |
+        echo "::group::Run Python tests - CPU"
+        python -m unittest discover python/tests -v
+        echo "::endgroup::"
 
     - name: Run Python tests - GPU
       if: inputs.run-tests == 'true'
@@ -40,22 +46,30 @@ runs:
       env:
         LOW_MEMORY: 1
         DEVICE: gpu
-      run: python -m tests discover python/tests -v
+      run: |
+        echo "::group::Run Python tests - GPU"
+        python -m tests discover python/tests -v
+        echo "::endgroup::"
 
     - name: Build CPP only
       if: inputs.build-type == 'debug'
       shell: bash
       run: |
+        echo "::group::Build CPP"
         cmake . -B build \
           -DMLX_BUILD_CUDA=ON \
           -DCMAKE_CUDA_COMPILER=${{ inputs.nvcc-location }} \
           -DCMAKE_BUILD_TYPE=DEBUG
         cmake --build build -j $(nproc)
+        echo "::endgroup::"
     
     - name: Run CPP tests
       if: ${{ inputs.build-type == 'debug' && inputs.run-tests == 'true' }}
       shell: bash
-      run: ./build/tests/tests -sfe="*fft_tests.cpp,*linalg_tests.cpp"
+      run: |
+        echo "::group::Run CPP Tests"
+        ./build/tests/tests -sfe="*fft_tests.cpp,*linalg_tests.cpp"
+        echo "::endgroup::"
 
     - name: Build Python package
       if: inputs.build-type == 'release'
diff --git a/.github/actions/build-docs/action.yml b/.github/actions/build-docs/action.yml
index 43af840815..042346cd6f 100644
--- a/.github/actions/build-docs/action.yml
+++ b/.github/actions/build-docs/action.yml
@@ -10,18 +10,22 @@ runs:
     - name: Install dependencies
       shell: sh
       run: |
+        echo "::group::Install dependencies"
         brew install doxygen
         uv pip install --upgrade pip cmake
         uv pip install -r docs/requirements.txt
         uv pip install . -v
+        echo "::endgroup::"
   
     - name: Build documentation
       shell: bash
       run: |
+        echo "::group::Build documentation"
         source .venv/bin/activate
         cd docs
         doxygen
         make html O=-W
+        echo "::endgroup::"
     
     - name: Create artifact tar
       shell: sh
diff --git a/.github/actions/build-linux/action.yml b/.github/actions/build-linux/action.yml
index 6a1c6b9820..3ee69dd7a3 100644
--- a/.github/actions/build-linux/action.yml
+++ b/.github/actions/build-linux/action.yml
@@ -28,18 +28,24 @@ runs:
       shell: sh
       env:
         CMAKE_ARGS: "-DCMAKE_COMPILE_WARNING_AS_ERROR=ON"
-      run: pip install -e ".[dev]" -v
+      run: |
+        echo "::group::Install Python Package"
+        pip install -e ".[dev]" -v
+        echo "::endgroup::"
     
     - name: Generate package stubs
       shell: sh
       run: |
+        echo "::group::Generate package stubs"
         pip install typing_extensions
         python setup.py generate_stubs
-    
+        echo "::endgroup::"
+
     - name: Run Python tests
       if: inputs.run-tests == 'true'
       shell: bash
       run: |
+        echo "::group::Run Python tests"
         python -m unittest discover python/tests -v
         mpirun --bind-to none --allow-run-as-root -host localhost:8 -np 8 python python/tests/mpi_test_distributed.py
         mlx.launch --verbose -n 8 python/tests/ring_test_distributed.py -v 2> >(tee -a stderr.log >&2)
@@ -48,24 +54,31 @@ runs:
           echo "Distributed ring test failed";
           exit 1;
         fi
-    
+        echo "::endgroup::"
+
     - name: Build CPP only
       if: inputs.build-type == 'debug'
       shell: bash
       run: |
+        echo "::group::Build CPP"
         mkdir -p build && cd build
         cmake .. -DMLX_BUILD_METAL=OFF -DCMAKE_BUILD_TYPE=DEBUG
         make -j $(nproc)
-    
+        echo "::endgroup::"
+
     - name: Run CPP tests
       if: ${{ inputs.build-type == 'debug' && inputs.run-tests == 'true' }}
       shell: sh
-      run: ./build/tests/tests
+      run: |
+        echo "::group::Run CPP Tests"
+        ./build/tests/tests
+        echo "::endgroup::"
     
     - name: Build Python package
       if: inputs.build-type == 'release'
       shell: bash
       run: |
+        echo "::group::Build Python Package"
         pip install auditwheel patchelf build
         python setup.py clean --all
         MLX_BUILD_STAGE=1 python -m build -w
@@ -76,3 +89,4 @@ runs:
         python setup.py clean --all
         MLX_BUILD_STAGE=2 python -m build -w
         auditwheel repair dist/mlx_cpu*.whl --plat manylinux_2_35_x86_64
+        echo "::endgroup::"
diff --git a/.github/actions/build-macos-release/action.yml b/.github/actions/build-macos-release/action.yml
index 7f5d13186d..a3579134b2 100644
--- a/.github/actions/build-macos-release/action.yml
+++ b/.github/actions/build-macos-release/action.yml
@@ -15,8 +15,10 @@ runs:
       env:
         MACOSX_DEPLOYMENT_TARGET: ${{ inputs.macos-target }}
       run: |
+        echo "::group::Build packages (target: $MACOSX_DEPLOYMENT_TARGET)"
         uv pip install build
         uv run --no-project setup.py clean --all
         MLX_BUILD_STAGE=1 uv run -m build -w
         uv run --no-project setup.py clean --all
         MLX_BUILD_STAGE=2 uv run -m build -w
+        echo "::endgroup::"
diff --git a/.github/actions/build-macos/action.yml b/.github/actions/build-macos/action.yml
index b1fe0876ee..3705e10da3 100644
--- a/.github/actions/build-macos/action.yml
+++ b/.github/actions/build-macos/action.yml
@@ -24,21 +24,27 @@ runs:
         DEBUG: 1
         DEV_RELEASE: 1
       run: |
+        echo "::group::Install dependencies"
         uv pip install --upgrade pip
         uv pip install cmake setuptools nanobind==2.4.0
         uv pip install -e . -v
+        echo "::endgroup::"
 
     - name: Generate package stubs
       shell: bash
       run: |
+        echo "::group::Generate package stubs"
         uv pip install typing_extensions
         uv run --no-project setup.py generate_stubs
+        echo "::endgroup::"
 
     - name: Install tests dependencies
       if: inputs.run-tests == 'true'
       shell: sh
       run: |
+        echo "::group::Install test dependencies"
         uv pip install numpy torch tensorflow unittest-xml-reporting
+        echo "::endgroup::"
 
     - name: Run Python tests
       if: inputs.run-tests == 'true'
@@ -46,43 +52,53 @@ runs:
       env:
         LOW_MEMORY: 1
       run: |
+        echo "::group::Run Python tests"
         DEVICE=cpu uv run -m xmlrunner discover -v python/tests -o test-results/cpu
         DEVICE=gpu METAL_DEVICE_WRAPPER_TYPE=1 METAL_DEBUG_ERROR_MODE=0 uv run -m xmlrunner discover -v python/tests -o test-results/gpu
         mpirun --bind-to none -host localhost:8 -np 8 -x DYLD_LIBRARY_PATH=/opt/homebrew/lib/ python python/tests/mpi_test_distributed.py
         mlx.launch --verbose -n 8 python/tests/ring_test_distributed.py -v 2> >(tee -a stderr.log >&2)
         if $(grep "\[WARN\]" stderr.log); then echo "Distributed ring test failed"; exit 1; fi
+        echo "::endgroup::"
     
     - name: Build example extension
       if: inputs.run-tests == 'true'
       shell: bash
       run: |
+        echo "::group::Build example extension"
         cd examples/extensions
         uv pip install -r requirements.txt
         uv run --no-project setup.py build_ext --inplace
         uv run --no-project test.py
-    
+        echo "::endgroup::"
+
     - name: Build CPP only
       if: inputs.run-tests == 'true'
       shell: bash
       run: |
+        echo "::group::Build CPP"
         mkdir -p build
         cd build
         cmake ..
         make -j $(sysctl -n hw.ncpu)
+        echo "::endgroup::"
     
-    - name: Run CPP tests
+    - name: Run CPP Tests
       if: inputs.run-tests == 'true'
       shell: bash
       env:
         DEVICE: gpu
         METAL_DEVICE_WRAPPER_TYPE: 1
         METAL_DEBUG_ERROR_MODE: 0
-      run: ./build/tests/tests
+      run: |
+        echo "::group::Run CPP Tests"
+        ./build/tests/tests
+        echo "::endgroup::"
     
     - name: Build small binary with JIT
       if: inputs.run-tests == 'true'
       shell: bash
       run: |
+        echo "::group::Build JIT Binary"
         mkdir -p build
         cd build
         cmake .. -DCMAKE_BUILD_TYPE=MinSizeRel \
@@ -92,7 +108,8 @@ runs:
           -DMLX_BUILD_GGUF=OFF \
           -DMLX_METAL_JIT=ON
         make -j $(sysctl -n hw.ncpu)
-    
+        echo "::endgroup::"
+
     - name: Run Python tests with JIT
       if: inputs.run-tests == 'true'
       shell: bash
@@ -102,11 +119,13 @@ runs:
         METAL_DEVICE_WRAPPER_TYPE: 1
         METAL_DEBUG_ERROR_MODE: 0
       run: |
+        echo "::group::Test JIT Binary"
         CMAKE_ARGS="-DMLX_METAL_JIT=ON" \
           uv pip install -e . -v
         uv run -m xmlrunner discover \
             -v python/tests \
             -o test-results/gpu_jit
+        echo "::endgroup::"
 
     - name: Build macOS 14 package
       if: inputs.build-type == 'release'
diff --git a/.github/actions/setup-linux/action.yml b/.github/actions/setup-linux/action.yml
index bcf3aa94d2..3c9eb46b80 100644
--- a/.github/actions/setup-linux/action.yml
+++ b/.github/actions/setup-linux/action.yml
@@ -28,9 +28,11 @@ runs:
         TZ: Etc/UTC
       shell: bash
       run: |
+        echo "::group::Install dependencies"
         sudo apt-get update
         sudo apt-get install -y libblas-dev liblapack-dev liblapacke-dev tzdata zip
         sudo apt autoremove -y
+        echo "::endgroup::"
 
     - uses: actions/setup-python@v6
       with: