Delete deprecated UintxLayout and related code

[jerryzh168]

Stack from ghstack (oldest at bottom):

• Delete deprecated TensorCoreTiledLayout and related code #4153
• Move bitpacking.py to prototype and add uintx_utils.py #4152
• Delete deprecated PlainLayout, PlainAQTTensorImpl and related v1 code paths #4151
• Delete deprecated PackedLinearInt8DynamicActivationIntxWeightLayout and related code #4150
• Delete deprecated Int4CPULayout and related code #4149
• Delete deprecated SemiSparseLayout and related code #4148
• Delete deprecated Int4XPULayout and related code #4147
• Delete deprecated QDQLayout and related code #4146
• -> Delete deprecated UintxLayout and related code #4145

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ You can merge normally! (1 Unrelated Failure)

As of commit ff2a214 with merge base fe98658 ():

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• Run Regression Tests / test-nightly (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch --index-url https://downloa... / linux-job (gh) (trunk failure)
  test/test_low_bit_optim.py::TestFSDP2::test_fsdp2

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

[jerryzh168]

@claude review

[claude]

Claude finished @jerryzh168's task in 2m 10s —— View job

---

Review of PR #4145: Delete deprecated UintxLayout and related code

• Gather context and review the diff
• Review deleted files for completeness
• Check remaining references to removed code
• Post review feedback

---

Summary

This PR removes the deprecated UintxLayout compatibility shim from torchao.dtypes.uintx.uintx_layout and the UintxAQTTensorImpl class, renames the implementation file from uintx_layout.py to uintx_utils.py, and updates import paths in the two remaining consumers (autoround and codebook). The test file for the deprecated path is also deleted.

Issue: Removing UintxAQTTensorImpl breaks autoround at runtime

This is the main concern with this PR.

The @register_layout(UintxLayout) decorator on UintxAQTTensorImpl was removed (the entire class was deleted from the renamed uintx_utils.py). However, torchao/prototype/autoround/core.py:203 still creates a UintxLayout and passes it as _layout to to_affine_quantized_intx_static():

# torchao/prototype/autoround/core.py:192-213
from torchao.prototype.dtypes.uintx.uintx_utils import (
    _BIT_WIDTH_TO_DTYPE,
    UintxLayout,
)
...
_layout = UintxLayout(dtype=dtype, pack_dim=pack_dim)
return to_affine_quantized_intx_static(
    input_float=input_float,
    ...
    _layout=_layout,
)

Inside from_hp_to_intx_static (torchao/dtypes/affine_quantized_tensor.py:462), this calls:

tensor_impl_ctr = get_tensor_impl_constructor(type(_layout))

Since UintxAQTTensorImpl (the class previously registered via @register_layout(UintxLayout)) has been deleted, this lookup will fail at runtime. The autoround code path that uses non-4-bit quantization will be broken.

Suggestion: Either:

1. Keep UintxAQTTensorImpl and @register_layout(UintxLayout) in uintx_utils.py (since autoround still depends on it), or
2. Migrate the autoround code to not use UintxLayout/AffineQuantizedTensor at all in this PR.

Fix this →

Minor observations

1. Import cleanup looks correct: The from torchao.dtypes.uintx.uintx_layout import UintxLayout removal from torchao/dtypes/__init__.py and torchao/dtypes/uintx/__init__.py is consistent — no stale references remain to the deleted file.

2. Codebook import update is fine: torchao/prototype/quantization/codebook/codebook_quantized_tensor.py:12 correctly updated its import from uintx_layout to uintx_utils and only uses _DTYPE_TO_BIT_WIDTH and UintxTensor (not UintxLayout), so it won't hit the register_layout issue.

3. Good documentation in new file: The docstring at the top of uintx_utils.py clearly explains these are retained utilities pending migration of autoround and codebook.

4. Duplicate dict initialization removed: The old code had an empty _DTYPE_TO_BIT_WIDTH = {} / _BIT_WIDTH_TO_DTYPE = {} followed by the real definitions — this is cleaned up in the new file. Good.

Verdict

The PR is clean aside from the runtime breakage for autoround caused by removing UintxAQTTensorImpl without updating or removing the UintxLayout usage in autoround/core.py. This needs to be addressed before merging.

---

[jerryzh168]

@claude review again

[claude]

Claude finished @jerryzh168's task in 1m 15s —— View job

---

Re-review of PR #4145: Delete deprecated UintxLayout and related code

• Gather context and review the current diff
• Check if previous review concerns were addressed
• Review remaining references to removed code
• Post review feedback

---

Summary

This PR removes the deprecated UintxLayout compatibility shim from torchao/dtypes/uintx/uintx_layout.py, renames the implementation file from uintx_layout.py to uintx_utils.py, cleans up the __init__.py re-exports, deletes the associated deprecation test, and updates import paths in autoround and codebook.

Previous concern resolved

The main issue from the prior review — that UintxAQTTensorImpl with @register_layout(UintxLayout) was being deleted, which would break autoround at runtime — has been fixed. UintxAQTTensorImpl is now retained in uintx_utils.py:223-237 with a clear TODO comment explaining it should be removed once autoround is migrated.

Current assessment

The PR looks good:

1. No stale references: No remaining imports from torchao.dtypes.uintx.uintx_layout or torchao.prototype.dtypes.uintx.uintx_layout exist in the codebase.

2. Import paths updated correctly: Both autoround/core.py:192 and codebook/codebook_quantized_tensor.py:12 now import from uintx_utils instead of uintx_layout.

3. Public API cleaned up properly: UintxLayout is removed from torchao/dtypes/__init__.py and torchao/dtypes/uintx/__init__.py. The prototype __init__.py files are emptied since nothing needs to be re-exported.

4. register_layout registration preserved: The @register_layout(UintxLayout) on UintxAQTTensorImpl is kept in uintx_utils.py, so the get_tensor_impl_constructor(type(_layout)) call in autoround will continue to work.

5. Clean code improvements: Duplicate empty dict initialization removed, docstrings simplified, good module-level docstring explaining the file's purpose and deprecation plan.

Verdict

LGTM. The previous concern is addressed and the PR is clean.

---

View job