Address compile InductorError in certain PyTorch versions

.github/workflows/check-links.yml

@@ -51,7 +51,7 @@ jobs:
         run: |
           set -euo pipefail
           source .venv/bin/activate
-          PYTHONPATH=.github/scripts pytest -p check_links_plugin --check-links ./ \
+          PYTHONPATH=.github/scripts pytest -p no:python -p check_links_plugin --check-links ./ \
             --check-links-ignore "https://platform.openai.com/*" \
             --check-links-ignore "https://openai.com/*" \
             --check-links-ignore "https://arena.lmsys.org" \

ch02/01_main-chapter-code/ch02_main.ipynb

@@ -523,7 +523,10 @@
     "        \n",
     "        if enable_tensor_cores:\n",
     "            major, minor = map(int, torch.__version__.split(\".\")[:2])\n",
-    "            if (major, minor) >= (2, 9):\n",
+    "            # PyTorch 2.9 and 2.10 still read the legacy TF32 setting in torch.compile.\n",
+    "            # See https://github.com/pytorch/pytorch/issues/166387\n",
+    "            # and https://github.com/rasbt/reasoning-from-scratch/issues/256\n",
+    "            if (major, minor) >= (2, 11):\n",
     "                torch.backends.cuda.matmul.fp32_precision = \"tf32\"\n",
     "                torch.backends.cudnn.conv.fp32_precision = \"tf32\"\n",
     "            else:\n",

pyproject.toml

@@ -61,6 +61,9 @@ extra = [
     "chainlit>=2.8.3",  # Appendix G, Chat UI
 ]
 
+[tool.ruff.lint]
+select = ["E4", "E7", "E9", "F"]
+
 [tool.setuptools.packages.find]
 where = ["."]
 include = ["reasoning_from_scratch", "reasoning_from_scratch.*"]

reasoning_from_scratch/ch02.py

@@ -14,7 +14,10 @@ def get_device(enable_tensor_cores=True):
 
         if enable_tensor_cores:
             major, minor = map(int, torch.__version__.split(".")[:2])
-            if (major, minor) >= (2, 9):
+            # PyTorch 2.9 and 2.10 still read the legacy TF32 setting in torch.compile.
+            # See https://github.com/pytorch/pytorch/issues/166387
+            # and https://github.com/rasbt/reasoning-from-scratch/issues/256
+            if (major, minor) >= (2, 11):
                 torch.backends.cuda.matmul.fp32_precision = "tf32"
                 torch.backends.cudnn.conv.fp32_precision = "tf32"
             else:

tests/test_ch02.py

@@ -2,6 +2,10 @@
 # Source for "Build a Reasoning Model (From Scratch)": https://mng.bz/lZ5B
 # Code repository: https://github.com/rasbt/reasoning-from-scratch
 
+import subprocess
+import sys
+
+import pytest
 import torch
 
 from reasoning_from_scratch.ch02 import (
@@ -52,6 +56,34 @@ def test_get_device_returns_torch_device(capsys):
     assert device.type in ("cpu", "cuda", "mps")
 
 
+def test_get_device_tf32_setting_is_compatible_with_torch_compile():
+    torch_version = tuple(map(int, torch.__version__.split(".")[:2]))
+    if not (2, 9) <= torch_version < (2, 11):
+        pytest.skip("PyTorch's mixed TF32 API bug affects versions 2.9 and 2.10")
+
+    # PyTorch 2.9 and 2.10 still read the legacy allow_tf32 flag inside
+    # torch.compile. Using the new API in get_device() triggers a mixed-API error.
+    code = """
+import torch
+
+from reasoning_from_scratch.ch02 import get_device
+
+torch.cuda.is_available = lambda: True
+get_device()
+torch.backends.cuda.matmul.allow_tf32
+"""
+    result = subprocess.run(
+        [sys.executable, "-c", code],
+        capture_output=True,
+        text=True,
+    )
+
+    assert result.returncode == 0, (
+        "get_device() selected a TF32 API that is incompatible with "
+        f"torch.compile:\n{result.stderr}"
+    )
+
+
 def test_generate_text_basic_stops_on_eos():
     # batch_size = 1
     # seq_len = 3