Unify get_block_size

[Xia-Weiwen]

Summary
This PR merges the two definitions of get_block_size in ao/torchao/quantization/pt2e/observer.py and torchao/quantization/observer.py and put it in torchao/quantization/utils.py
It also removes duplicate definitions of granularities in the pt2e path.

Test plan
This util function is used in many places so it should have been covered by all kinds of UT.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 f07ab47 with merge base eadead5 ():
💚 Looks good so far! There are no failures yet. 💚

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

[Copilot]

Pull Request Overview

This PR consolidates the get_block_size function by moving it from two separate locations (torchao/quantization/pt2e/observer.py and torchao/quantization/observer.py) to a centralized location in torchao/quantization/utils.py. The PR also removes duplicate granularity class definitions from the pt2e path, centralizing granularity-related functionality.

• Moves get_block_size function to a central utils.py location
• Removes duplicate granularity class definitions from pt2e/observer.py
• Updates import statements across multiple files to use the centralized functions

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/quantization/utils.py	Adds the unified get_block_size function and imports granularity classes
torchao/quantization/quantize_/workflows/float8/float8_tensor.py	Updates import to use get_block_size from utils instead of observer
torchao/quantization/quant_api.py	Updates import to use get_block_size from utils instead of observer
torchao/quantization/qat/fake_quantizer.py	Updates import to use get_block_size from utils instead of observer
torchao/quantization/pt2e/observer.py	Removes duplicate granularity classes and get_block_size function
torchao/quantization/pt2e/_affine_quantization.py	Updates import to use get_block_size from utils
torchao/quantization/pt2e/init.py	Removes exports for granularity classes and get_block_size
torchao/quantization/observer.py	Removes original get_block_size function and unused import
torchao/quantization/init.py	Adds export for Granularity base class
test/quantization/pt2e/test_quantize_pt2e.py	Updates imports to use granularity classes from main quantization module

---

_{Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.}

[Copilot]

The unified get_block_size function has inconsistent behavior for PerGroup compared to the original pt2e version. The original pt2e version enforced a 2D input shape constraint and returned (1, granularity.group_size), while this version allows any dimensional input and returns (1,) * (len(input_shape) - 1) + (granularity.group_size,). This could break existing code that depends on the 2D constraint.

Suggested change

	assert input_shape[-1] % granularity.group_size == 0, (
	f"Group size {granularity.group_size} does not divide input shape {input_shape}"
	)
	return (1,) * (len(input_shape) - 1) + (granularity.group_size,)
	assert len(input_shape) == 2, (
	f"PerGroup granularity requires 2D input shape, got {input_shape}"
	)
	assert input_shape[-1] % granularity.group_size == 0, (
	f"Group size {granularity.group_size} does not divide input shape {input_shape}"
	)
	return (1, granularity.group_size)

[Xia-Weiwen]

We should allow any dimensional input instead of requiring 2D.

[Xia-Weiwen]

Hi @jerryzh168 Could you please review this PR? Thanks.

[jerryzh168]

@Xia-Weiwen thanks, can you fix ruff?

[jerryzh168]

looks good, please make sure CI passes

[cccclai]

How about PerBlock?

[jerryzh168]

sorry reverting this first, I think we should keep BC

[Xia-Weiwen]

How about PerBlock?

PerBlock is added in the reopened PR #3059. However, you can easily get block size by granularity.block_size if granularity is PerBlock.