Skip to content

Add squeeze/unsqueeze support to IntQuantTensor - #1622

Draft
mumallaeng wants to merge 1 commit into
Xilinx:masterfrom
mumallaeng:quanttensor-squeeze-unsqueeze
Draft

mumallaeng wants to merge 1 commit into
Xilinx:masterfrom
mumallaeng:quanttensor-squeeze-unsqueeze

Conversation

@mumallaeng

@mumallaeng mumallaeng commented Sep 20, 2026 •

Copy link
Copy Markdown

Reason for this PR

QuantTensor had no squeeze/unsqueeze, so torch.squeeze/torch.unsqueeze (and the .squeeze()/.unsqueeze() methods) either raised or silently dropped quantization metadata (scale/zero_point/bit_width) for IntQuantTensor.

Closes #891.

Changes Made in this PR

Added squeeze/unsqueeze to IntMixin in base_quant_tensor.py, following the same pattern already used by transpose/permute: per-channel metadata (same rank as the value tensor) is squeezed/unsqueezed alongside the value, per-tensor (lower-rank) metadata is left untouched. Registered torch.squeeze/torch.unsqueeze handlers in torch_handler.py so the free-function form dispatches correctly too.

FloatQuantTensor is intentionally left out of scope, per the prior PR #941 discussion - a follow-up can extend this once FloatQuantTensor's metadata shape handling is worked out separately.

Testing Summary

Added test_quant_tensor_squeeze and test_quant_tensor_unsqueeze to tests/brevitas/quant_tensor/test_quant_tensor.py, covering both per-tensor and per-channel scale. Ran the full test_quant_tensor.py file locally (30/30 pass). Confirmed the new tests fail with AttributeError against the pre-change source and pass after. isort/yapf (pinned versions from .pre-commit-config.yaml) show no diff on the changed files.

Risk Highlight

  • This PR includes code from another work (please detail).
  • This PR contains API-breaking changes.
  • This PR depends on work in another PR (please provide links/details).
  • This PR introduces new dependencies (please detail).
  • There are coverage gaps not covered by tests.
  • Documentation updates required in subsequent PR.

Checklist

  • Code comments added to any hard-to-understand areas, if applicable.
  • Changes generate no new warnings.
  • Updated any relevant tests, if applicable.
  • No conflicts with destination master branch.
  • I reviewed my own code changes.
  • Initial CI/CD passing.
  • 1+ reviews given, and any review issues addressed and approved.
  • Post-review full CI/CD passing.

QuantTensor had no squeeze/unsqueeze, so torch.squeeze/torch.unsqueeze
(and the .squeeze()/.unsqueeze() methods) either raised or silently
dropped quantization metadata (scale/zero_point/bit_width) for
IntQuantTensor.

Add squeeze/unsqueeze to IntMixin following the same pattern already
used by transpose/permute: per-channel metadata (same rank as the
value tensor) is squeezed/unsqueezed alongside the value, per-tensor
(lower-rank) metadata is left untouched. Register torch.squeeze/
torch.unsqueeze handlers in torch_handler.py so the free-function form
dispatches correctly too.

FloatQuantTensor is intentionally left out of scope, per the prior
PR Xilinx#941 discussion - a follow-up can extend this once FloatQuantTensor's
metadata shape handling is worked out separately.

Signed-off-by: mumallaeng <mumallaeng@outlook.com>
@Giuseppe5

Copy link
Copy Markdown
Collaborator

Hey, thanks for opening this PR. Unless this is causing some issues for you, we might want to skip this change for now.

We are planning major changes around QuantTensor #1579, and we might want to merge that before applying other changes to QuantTensor.

@mumallaeng

Copy link
Copy Markdown
Author

@Giuseppe5 Thanks for the heads-up, that makes sense. I'll convert this to a draft and hold it until #1579 lands, then re-check whether squeeze/unsqueeze is still needed on top of the tensor-subclass QuantTensor.

@mumallaeng
mumallaeng marked this pull request as draft September 22, 2026 01:00

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Add squeeze / unsqueeze operations to quant invariant functions in torch_handler.py

2 participants