Remove check_cpu_version and check_xpu_version helpers

[Xia-Weiwen]

This PR removes check_cpu/xpu_version since they are ambiguous and changes all uses of these helpers to _is_device("cpu/xpu", tensor.device).

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/4211

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit c44ba5c with merge base 42ff8ed ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[jerryzh168]

@claude review

[claude]

Claude finished @jerryzh168's task in 4m 9s —— View job

---

Review of PR #4211: Remove check_cpu_version and check_xpu_version helpers

• Gather context and read changed files
• Analyze changes and provide review feedback

---

Overall: This is a good cleanup. Removing the version-gated helpers makes sense since the minimum supported torch version is now 2.8.0, which is above both the check_cpu_version default (2.6.0) and check_xpu_version default (2.8.0). The version checks were always true and therefore dead code.

No remaining references to the removed functions exist in the codebase.

One issue to flag:

The replacement comparisons use device == "cpu" and device == "xpu" where device is a torch.device object (e.g. from tensor.device). In PyTorch, torch.device does support __eq__ with strings, so torch.device("cpu") == "cpu" returns True. This works today.

However, note that this codebase is inconsistent — some places use .device == "cpu" while others use .device.type == "cpu" (e.g., torchao/kernel/intmm.py:45, torchao/quantization/quantize_/workflows/int4/int4_plain_int32_tensor.py:92). There's also an existing is_device() helper in torchao/dtypes/utils.py:94 specifically for this purpose. This inconsistency isn't introduced by this PR (the old check_cpu_version was doing .type comparison internally), but it's worth being aware of. Not a blocker.

Specific observations per file:

1. torchao/utils.py — Clean removal of both functions. No exports to update.

2. test/dtypes/test_affine_quantized.py — The removed block (if do_int4: if check_cpu_version(device): pass elif check_xpu_version(device): pass) was a no-op (all branches were pass), so removing it is correct and simplifies the code. Note that the do_int4 parameter is now completely unused in get_quantization_functions — it's still accepted but never read. You may want to clean that up in a follow-up.

3. test/quantization/test_quant_primitives.py, torchao/quantization/utils.py, torchao/kernel/intmm.py, torchao/prototype/hqq/hqq_tinygemm_linear.py — Straightforward replacements of check_cpu_version(device) → device == "cpu" and check_xpu_version(device) → device == "xpu". The negated forms not (check_cpu_version(...)) correctly became != "cpu". Logic is preserved.

Summary: LGTM. The changes are correct and CI is green. The only minor follow-up suggestion is removing the now-unused do_int4 and int4_zp_int parameters from get_quantization_functions in test/dtypes/test_affine_quantized.py.

---

[jerryzh168]

this condition seems complicated, might be good to simplify if possible

[Xia-Weiwen]

Oh, I am not familiar with the code here. And I feel the condition cannot be simplified. However, we can probably add a comment here to clarify the condition.
CC the last author @yanbing-j

[jerryzh168]

There's also an existing is_device() helper in torchao/dtypes/utils.py:94 specifically for this purpose.

I feel we should move this helper to torchao/utils.py and use this helper everywhere

[Xia-Weiwen]

There's also an existing is_device() helper in torchao/dtypes/utils.py:94 specifically for this purpose.

I feel we should move this helper to torchao/utils.py and use this helper everywhere

Hi @jerryzh168 I feel that util function does not really simplify code.

is_device("cpu", tensor.device)
tensor.device.type == "cpu"

not is_device("cpu", tensor.device)
tensor.device.type != "cpu"

However, if you think the util function is better, I will use that.

[Xia-Weiwen]

There's also an existing is_device() helper in torchao/dtypes/utils.py:94 specifically for this purpose.

I feel we should move this helper to torchao/utils.py and use this helper everywhere

Hi @jerryzh168 I feel that util function does not really simplify code.

is_device("cpu", tensor.device)
tensor.device.type == "cpu"

not is_device("cpu", tensor.device)
tensor.device.type != "cpu"

However, if you think the util function is better, I will use that.

@jerryzh168 I have removed the is_device utils and added is_on_device(tensor, device) and not_on_device(tensor, devices) in torchao.utils. These two utils might looks less confusing and more convenient. How does that sound to you? Thanks.

[Xia-Weiwen]

@claude Please review again

[claude]

Claude finished @Xia-Weiwen's task in 1m 47s —— View job

---

Re-review of PR #4211: Remove check_cpu_version and check_xpu_version helpers

• Gather context and read changed files
• Analyze new is_on_device and not_on_device helpers
• Review all call sites for correctness
• Provide review feedback

---

Overall: This revision is a nice improvement over the previous version. The new is_on_device / not_on_device helpers in torchao/utils.py are clean, and all call sites have been updated consistently. No remaining references to check_cpu_version, check_xpu_version, or is_device from torchao/dtypes/utils.py exist in the codebase.

New helpers (torchao/utils.py:1279-1288):

The API is straightforward and well-designed:

• is_on_device(tensor, "cpu") — reads naturally
• not_on_device(tensor, ["cpu", "xpu"]) — the list form is a nice touch for the multi-device exclusion pattern that appears several times

Two minor observations:

1. __all__ not updated — The new is_on_device and not_on_device functions are not added to __all__ in torchao/utils.py:20-39. The old check_cpu_version and check_xpu_version weren't in __all__ either, so this is consistent with the prior state. But since these are public helpers imported across the codebase, it may be worth adding them. Not a blocker.

2. Inconsistency at linear_quant_modules.py:119 — One site uses device.type == "cpu" directly instead of is_on_device. This is because device here is a constructor parameter (not a tensor), so is_on_device doesn't apply. This is fine — the helper is tensor-specific by design, and using .type directly on a torch.device is the right call here.

Per-file observations:

File	Status
torchao/utils.py	Clean removal of old helpers, clean addition of new ones
torchao/dtypes/utils.py	is_device removed, unused Union import cleaned up
torchao/kernel/intmm.py	Correct replacement
torchao/quantization/utils.py	Correct replacements. The condition at line 483-485 is still complex but this PR doesn't change its logic, just the device-check syntax
torchao/quantization/linear_quant_modules.py	Correct. Line 119 uses device.type == "cpu" directly since device is not a tensor — this is appropriate
torchao/quantization/qat/linear.py	Correct replacement
torchao/prototype/hqq/hqq_tinygemm_linear.py	Correct. Previously mixed is_device and check_cpu_version; now uniformly uses is_on_device
torchao/prototype/tensor_conversion/api.py	Correct replacement, removed stale # TODO: move the function to torchao.utils comment
test/dtypes/test_affine_quantized.py	Dead code removed (all branches were pass). Note: do_int4 param is now unused — could be cleaned up in a follow-up
test/quantization/test_quant_primitives.py	Correct replacements

CI: Ruff linting and PR label check pass. GPU/regression tests still running.

Summary: LGTM. The changes are correct and the new helpers are an improvement in readability and consistency. The only follow-up suggestions (not blockers) are adding the new functions to __all__ and cleaning up the unused do_int4 parameter in tests.

--- |

[Copilot]

Pull request overview

This PR removes the check_cpu_version / check_xpu_version helpers (and related call sites) and consolidates device checks into new is_on_device / not_on_device utilities.

Changes:

• Delete check_cpu_version and check_xpu_version from torchao/utils.py and introduce is_on_device / not_on_device.
• Replace device/version helper usage across quantization, kernel, and prototype code with the new device-check utilities.
• Remove the now-unused is_device helper from torchao/dtypes/utils.py and update tests accordingly.

Reviewed changes

Copilot reviewed 10 out of 10 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
torchao/utils.py	Removes CPU/XPU version-check helpers; adds shared tensor device predicates.
torchao/quantization/utils.py	Updates int4 packing/unpacking conditionals to use is_on_device / not_on_device.
torchao/quantization/qat/linear.py	Replaces is_device(...) usage with is_on_device(...).
torchao/quantization/linear_quant_modules.py	Replaces is_device(...) usage with is_on_device(...) and direct device.type checks.
torchao/prototype/tensor_conversion/api.py	Switches CUDA device checks to is_on_device.
torchao/prototype/hqq/hqq_tinygemm_linear.py	Removes CPU version helper usage in favor of is_on_device.
torchao/kernel/intmm.py	Replaces CPU-version helper usage with is_on_device.
torchao/dtypes/utils.py	Removes the unused is_device helper.
test/quantization/test_quant_primitives.py	Updates test logic to use is_on_device / not_on_device.
test/dtypes/test_affine_quantized.py	Removes now-unused imports and dead conditional block related to CPU/XPU version helpers.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[jerryzh168]

There's also an existing is_device() helper in torchao/dtypes/utils.py:94 specifically for this purpose.

I feel we should move this helper to torchao/utils.py and use this helper everywhere

Hi @jerryzh168 I feel that util function does not really simplify code.

is_device("cpu", tensor.device)
tensor.device.type == "cpu"

not is_device("cpu", tensor.device)
tensor.device.type != "cpu"

However, if you think the util function is better, I will use that.

the util is to make sure it works for multiple situations, like device being string, being torch.device type etc., it's more for robustness of the check

I don't understand the motivation of is_on_device and is_not_on_device, why is_device is not enough here?

[jerryzh168]

I see, I don't think we need special util just for having a list of args, I think it's better to just use is_device

[Xia-Weiwen]

There's also an existing is_device() helper in torchao/dtypes/utils.py:94 specifically for this purpose.

I feel we should move this helper to torchao/utils.py and use this helper everywhere

Hi @jerryzh168 I feel that util function does not really simplify code.

is_device("cpu", tensor.device)
tensor.device.type == "cpu"

not is_device("cpu", tensor.device)
tensor.device.type != "cpu"

However, if you think the util function is better, I will use that.

the util is to make sure it works for multiple situations, like device being string, being torch.device type etc., it's more for robustness of the check

I don't understand the motivation of is_on_device and is_not_on_device, why is_device is not enough here?

The motivation is to simplify code and make it more straightforward: the new utils work on tensors so that callers don't have to write tensor.device or tensor.deice.type, just tensor.

is_device("cpu", tensor.device)
tensor.device.type == "cpu"
is_on_device(tensor, "cpu")

not is_device("cpu", tensor.device) and not is_device("xpu", tensor.device)
tensor.device.type != "cpu" and tensor.device.type != "xpu"
not_on_device(tensor, ["cpu", "xpu"])

Anyway, I will rollback to is_device if you don't like them :)

[Xia-Weiwen]

@claude review

[claude]

Claude finished @Xia-Weiwen's task in 1m 53s —— View job

---

Review of PR #4211: Remove check_cpu_version and check_xpu_version helpers

• Gather context and read changed files
• Analyze changes for correctness
• Provide review feedback

---

Overall: This is a clean, focused PR. The check_cpu_version and check_xpu_version helpers combined a device type check with a torch version check, but since the minimum supported torch version is now 2.8.0 (above both the check_cpu_version default of 2.6.0 and check_xpu_version default of 2.8.0), the version checks were always true. Removing them and replacing with _is_device calls is the right cleanup.

CI: 6/7 checks passed, 1 still in progress (Regression Tests). Ruff linting, H100, L4, MPS, aarch64, and label check all pass.

Per-file observations:

File	Status
torchao/utils.py	Clean removal of check_cpu_version and check_xpu_version (12 lines deleted). No other changes needed — _is_device already exists at line 48.
torchao/kernel/intmm.py	Correct: check_cpu_version(scales1.device) → _is_device("cpu", scales1.device). Import updated.
torchao/prototype/hqq/hqq_tinygemm_linear.py	Correct: Two check_cpu_version calls and one _is_device import consolidated. Now imports only _is_device from torchao.utils.
torchao/quantization/utils.py	Correct: All check_cpu_version/check_xpu_version calls replaced with _is_device. The complex condition at lines 462-467 preserves the original logic exactly.
test/quantization/test_quant_primitives.py	Correct: Same pattern of replacements. Imports updated.

No remaining references to check_cpu_version or check_xpu_version exist anywhere in the codebase.

One observation (not a blocker): The argument order of _is_device is inconsistent across the codebase. This PR correctly uses _is_device("cpu", device) matching the function signature _is_device(target_device_str, device). However, some pre-existing callers in torchao/quantization/linear_quant_modules.py and torchao/quantization/qat/linear.py (not changed by this PR) use the reversed order: _is_device(q_weight.device.type, "cpu"). This happens to work because string equality is symmetric, but it's confusing. Worth a follow-up cleanup to standardize.

Summary: LGTM. The changes are correct, minimal, and CI looks good.

---

[jerryzh168]

yeah I think this is cleaner, please update the summary as well