ops_python: add image support (#3356)

* ops_python: add image support

* uops tests in their own CI

* fix ci

.github/workflows/test.yml

@@ -11,6 +11,32 @@ on:
   workflow_dispatch:
 
 jobs:
+  uops:
+    name: uops tests
+    runs-on: ubuntu-latest
+    timeout-minutes: 20
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v4
+    - name: Set up Python 3.12
+      uses: actions/setup-python@v5
+      with:
+        python-version: 3.12
+    - name: Cache python packages
+      uses: actions/cache@v3
+      with:
+        path: ${{ env.Python3_ROOT_DIR }}/lib/python3.8/site-packages
+        key: uops-packages-${{ hashFiles('**/setup.py') }}-3.12
+    - name: Install dependencies
+      run: pip install -e '.[testing]' --extra-index-url https://download.pytorch.org/whl/cpu
+    - name: Test IMAGE=2 support
+      run: |
+        IMAGE=2 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm
+      # TODO: this is broken
+      #IMAGE=2 PYTHON=1 python3 test/test_ops.py TestOps.test_simple_conv2d
+    - name: Test emulated METAL tensor cores
+      run: DEBUG=2 EMULATE_METAL=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_big_gemm
+
   linter:
     name: Linters
     runs-on: ubuntu-latest
@@ -55,8 +81,6 @@ jobs:
         PYTHONPATH="." python test/external/fuzz_shapetracker_math.py
     - name: Test shapetracker to_movement_ops
       run: PYTHONPATH="." python extra/to_movement_ops.py
-    - name: Test emulated METAL tensor cores
-      run: DEBUG=2 EMULATE_METAL=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_big_gemm
     - name: Use as an external package
       run: |
         mkdir $HOME/test_external_dir