test: add tracer-leak check

Signed-off-by: nstarman <[email protected]>

Author: nstarman

.github/workflows/run_tests.yml

@@ -32,3 +32,26 @@ jobs:
         run: |
           python -m pip install .
           python -m test
+
+  # Run a test with JAX tracer leak detection enabled
+  run-test-tracer:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v2
+
+      - name: Set up Python 3.13
+        uses: actions/setup-python@v2
+        with:
+          python-version: "3.13"
+
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          python -m pip install '.[tests]'
+
+      - name: Test tracer functionality
+        run: |
+          python -m pip install .
+          export JAX_CHECK_TRACER_LEAKS=1
+          python -m test --tracer